prompt.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import pandas as pd\n",
    "from tqdm import tqdm\n",
    "import torch\n",
    "import jieba\n",
    "import codecs\n",
    "import json\n",
    "from sklearn.utils import shuffle\n",
    "from typing import List, Dict\n",
    "from datasets import Dataset\n",
    "from sklearn.metrics import accuracy_score, precision_recall_fscore_support\n",
    "import numpy as np\n",
    "from transformers import (\n",
    "    Trainer,\n",
    "    BertTokenizer,\n",
    "    TrainingArguments,\n",
    "    BertForMaskedLM,\n",
    "    EarlyStoppingCallback\n",
    ")\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"1\"\n",
    "\n",
    "source_data_path = \"/data/christmas.wang/project/classification_base_project/data/origin/few_shot/\"\n",
    "bert_name_path = \"/data/tmp/christmas.wang/chinese_wwm_ext_pytorch\"\n",
    "\n",
    "label_2_id = {\"好\":1, \"差\":0}\n",
    "id_2_label = {\"1\":\"好\", \"0\":\"差\"}\n",
    "\n",
    "data_dict = {\n",
    "    \"train\": os.path.join(source_data_path, \"hotel_review_few_shot_train.csv\"),\n",
    "    \"test\": os.path.join(source_data_path, \"hotel_review_few_shot_test.csv\")\n",
    "}\n",
    "\n",
    "# f->few shot; z->zero_shot\n",
    "z_or_f = \"f\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Didn't find file /data/tmp/christmas.wang/chinese_wwm_ext_pytorch/added_tokens.json. We won't load it.\n",
      "Didn't find file /data/tmp/christmas.wang/chinese_wwm_ext_pytorch/special_tokens_map.json. We won't load it.\n",
      "Didn't find file /data/tmp/christmas.wang/chinese_wwm_ext_pytorch/tokenizer_config.json. We won't load it.\n",
      "Didn't find file /data/tmp/christmas.wang/chinese_wwm_ext_pytorch/tokenizer.json. We won't load it.\n",
      "loading file /data/tmp/christmas.wang/chinese_wwm_ext_pytorch/vocab.txt\n",
      "loading file None\n",
      "loading file None\n",
      "loading file None\n",
      "loading file None\n",
      "loading configuration file /data/tmp/christmas.wang/chinese_wwm_ext_pytorch/config.json\n",
      "loading configuration file /data/tmp/christmas.wang/chinese_wwm_ext_pytorch/config.json\n",
      "Model config BertConfig {\n",
      "  \"attention_probs_dropout_prob\": 0.1,\n",
      "  \"classifier_dropout\": null,\n",
      "  \"directionality\": \"bidi\",\n",
      "  \"gradient_checkpointing\": false,\n",
      "  \"hidden_act\": \"gelu\",\n",
      "  \"hidden_dropout_prob\": 0.1,\n",
      "  \"hidden_size\": 768,\n",
      "  \"initializer_range\": 0.02,\n",
      "  \"intermediate_size\": 3072,\n",
      "  \"layer_norm_eps\": 1e-12,\n",
      "  \"max_position_embeddings\": 512,\n",
      "  \"model_type\": \"bert\",\n",
      "  \"num_attention_heads\": 12,\n",
      "  \"num_hidden_layers\": 12,\n",
      "  \"pad_token_id\": 0,\n",
      "  \"pooler_fc_size\": 768,\n",
      "  \"pooler_num_attention_heads\": 12,\n",
      "  \"pooler_num_fc_layers\": 3,\n",
      "  \"pooler_size_per_head\": 128,\n",
      "  \"pooler_type\": \"first_token_transform\",\n",
      "  \"position_embedding_type\": \"absolute\",\n",
      "  \"transformers_version\": \"4.10.2\",\n",
      "  \"type_vocab_size\": 2,\n",
      "  \"use_cache\": true,\n",
      "  \"vocab_size\": 21128\n",
      "}\n",
      "\n",
      "loading weights file /data/tmp/christmas.wang/chinese_wwm_ext_pytorch/pytorch_model.bin\n",
      "Some weights of the model checkpoint at /data/tmp/christmas.wang/chinese_wwm_ext_pytorch were not used when initializing BertForMaskedLM: ['cls.seq_relationship.weight', 'cls.seq_relationship.bias']\n",
      "- This IS expected if you are initializing BertForMaskedLM 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 BertForMaskedLM from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "All the weights of BertForMaskedLM were initialized from the model checkpoint at /data/tmp/christmas.wang/chinese_wwm_ext_pytorch.\n",
      "If your task is similar to the task the model of the checkpoint was trained on, you can already use BertForMaskedLM for predictions without further training.\n"
     ]
    }
   ],
   "source": [
    "tokenizer = BertTokenizer.from_pretrained(bert_name_path)\n",
    "model = BertForMaskedLM.from_pretrained(bert_name_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_metrics(pred):\n",
    "    labels = pred.label_ids[:, 3]\n",
    "    preds = pred.predictions[:, 3].argmax(-1)\n",
    "    precision, recall, f1, _ = precision_recall_fscore_support(labels, preds, average='weighted')\n",
    "    acc = accuracy_score(labels, preds)\n",
    "    return {'accuracy': acc, 'f1': f1, 'precision': precision, 'recall': recall}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Unnamed: 0</th>\n",
       "      <th>label</th>\n",
       "      <th>review</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>5322</td>\n",
       "      <td>0</td>\n",
       "      <td>标准间太差房间还不如3星的而且设施非常陈旧.建议酒店把老的标准间从新改善.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>5323</td>\n",
       "      <td>0</td>\n",
       "      <td>服务态度极其差，前台接待好象没有受过培训，连基本的礼貌都不懂，竟然同时接待几个客人；大堂副理...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>5324</td>\n",
       "      <td>0</td>\n",
       "      <td>地理位置还不错，到哪里都比较方便，但是服务不象是豪生集团管理的，比较差。下午睡了一觉并洗了一...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>5325</td>\n",
       "      <td>0</td>\n",
       "      <td>1。我住的是靠马路的标准间。房间内设施简陋，并且的房间玻璃窗户外还有一层幕墙玻璃，而且不能打...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>5326</td>\n",
       "      <td>0</td>\n",
       "      <td>我这次是第5次住在长春的雁鸣湖大酒店。昨晚夜里停电。深夜我睡着了。我的钱包被内贼进入我的房间...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>5327</td>\n",
       "      <td>0</td>\n",
       "      <td>前台checkin花了20分钟，checkout25分钟，这是服务态度和没有做到位。信用卡刷...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>5328</td>\n",
       "      <td>0</td>\n",
       "      <td>有或者很少房!梯部不吸,但是有一些吸者仍然有服!我是不抽的人,成二手的受害者!(中13人口中...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>5329</td>\n",
       "      <td>0</td>\n",
       "      <td>酒店服务态度极差，设施很差，建议还是不要到那儿去。</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>5330</td>\n",
       "      <td>0</td>\n",
       "      <td>我3.6预定好的180的标间，当我到的时候竟然说有会议房间满了，我订的房间没有了，太不讲信誉...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>5331</td>\n",
       "      <td>0</td>\n",
       "      <td>房间的环境非常差,而且房间还不隔音，住的不舒服。</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>距离川沙公路较近,但是公交指示不对,如果是\"蔡陆线\"的话,会非常麻烦.建议用别的路线.房间较...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11</th>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>商务大床房，房间很大，床有2M宽，整体感觉经济实惠不错!</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>12</th>\n",
       "      <td>2</td>\n",
       "      <td>1</td>\n",
       "      <td>早餐太差，无论去多少人，那边也不加食品的。酒店应该重视一下这个问题了。房间本身很好。</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>13</th>\n",
       "      <td>3</td>\n",
       "      <td>1</td>\n",
       "      <td>宾馆在小街道上，不大好找，但还好北京热心同胞很多~宾馆设施跟介绍的差不多，房间很小，确实挺小...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14</th>\n",
       "      <td>4</td>\n",
       "      <td>1</td>\n",
       "      <td>CBD中心,周围没什么店铺,说5星有点勉强.不知道为什么卫生间没有电吹风</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>15</th>\n",
       "      <td>5</td>\n",
       "      <td>1</td>\n",
       "      <td>总的来说，这样的酒店配这样的价格还算可以，希望他赶快装修，给我的客人留些好的印象</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>16</th>\n",
       "      <td>6</td>\n",
       "      <td>1</td>\n",
       "      <td>价格比比较不错的酒店。这次免费升级了，感谢前台服务员。房子还好，地毯是新的，比上次的好些。早...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>17</th>\n",
       "      <td>7</td>\n",
       "      <td>1</td>\n",
       "      <td>不错，在同等档次酒店中应该是值得推荐的！</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>18</th>\n",
       "      <td>8</td>\n",
       "      <td>1</td>\n",
       "      <td>入住丽晶，感觉很好。因为是新酒店，的确有淡淡的油漆味，房间内较新。房间大小合适，卫生间设备齐...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>19</th>\n",
       "      <td>9</td>\n",
       "      <td>1</td>\n",
       "      <td>1。酒店比较新，装潢和设施还不错，只是房间有些油漆味。2。早餐还可以，只是品种不是很多。3。...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "    Unnamed: 0  label                                             review\n",
       "0         5322      0              标准间太差房间还不如3星的而且设施非常陈旧.建议酒店把老的标准间从新改善.\n",
       "1         5323      0  服务态度极其差，前台接待好象没有受过培训，连基本的礼貌都不懂，竟然同时接待几个客人；大堂副理...\n",
       "2         5324      0  地理位置还不错，到哪里都比较方便，但是服务不象是豪生集团管理的，比较差。下午睡了一觉并洗了一...\n",
       "3         5325      0  1。我住的是靠马路的标准间。房间内设施简陋，并且的房间玻璃窗户外还有一层幕墙玻璃，而且不能打...\n",
       "4         5326      0  我这次是第5次住在长春的雁鸣湖大酒店。昨晚夜里停电。深夜我睡着了。我的钱包被内贼进入我的房间...\n",
       "5         5327      0  前台checkin花了20分钟，checkout25分钟，这是服务态度和没有做到位。信用卡刷...\n",
       "6         5328      0  有或者很少房!梯部不吸,但是有一些吸者仍然有服!我是不抽的人,成二手的受害者!(中13人口中...\n",
       "7         5329      0                          酒店服务态度极差，设施很差，建议还是不要到那儿去。\n",
       "8         5330      0  我3.6预定好的180的标间，当我到的时候竟然说有会议房间满了，我订的房间没有了，太不讲信誉...\n",
       "9         5331      0                           房间的环境非常差,而且房间还不隔音，住的不舒服。\n",
       "10           0      1  距离川沙公路较近,但是公交指示不对,如果是\"蔡陆线\"的话,会非常麻烦.建议用别的路线.房间较...\n",
       "11           1      1                       商务大床房，房间很大，床有2M宽，整体感觉经济实惠不错!\n",
       "12           2      1         早餐太差，无论去多少人，那边也不加食品的。酒店应该重视一下这个问题了。房间本身很好。\n",
       "13           3      1  宾馆在小街道上，不大好找，但还好北京热心同胞很多~宾馆设施跟介绍的差不多，房间很小，确实挺小...\n",
       "14           4      1               CBD中心,周围没什么店铺,说5星有点勉强.不知道为什么卫生间没有电吹风\n",
       "15           5      1           总的来说，这样的酒店配这样的价格还算可以，希望他赶快装修，给我的客人留些好的印象\n",
       "16           6      1  价格比比较不错的酒店。这次免费升级了，感谢前台服务员。房子还好，地毯是新的，比上次的好些。早...\n",
       "17           7      1                               不错，在同等档次酒店中应该是值得推荐的！\n",
       "18           8      1  入住丽晶，感觉很好。因为是新酒店，的确有淡淡的油漆味，房间内较新。房间大小合适，卫生间设备齐...\n",
       "19           9      1  1。酒店比较新，装潢和设施还不错，只是房间有些油漆味。2。早餐还可以，只是品种不是很多。3。..."
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_train = pd.read_csv(data_dict[\"train\"], encoding=\"utf-8\")\n",
    "df_test = pd.read_csv(data_dict[\"test\"], encoding=\"utf-8\")\n",
    "df_train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████| 20/20 [00:00<00:00, 470.79it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1355\n",
      "['地理位置还不错，到哪里都比较方便，但是服务不象是豪生集团管理的，比较差。下午睡了一觉并洗了一个澡，本来想让酒店再来打扫，酒店[MASK]，一下，所以，打开了，请打扫的服务灯，可是到晚上回酒店，发现打扫得服务灯被关掉了，而房间还是没有打扫过。', '我这次是第5次住在长春的雁鸣湖大酒店。昨晚夜里停电。深夜我睡着了。我的钱包被内贼进入我的房间，偷了我近1000元和4张信用卡。。。我的证件和外币，数码相机等都在房间的保险箱里，原封不动。我打了好几个小时的长途电话来处理我的信用卡的冻结。我报案了，这个4星酒店的保安摄像探头竟然坏了，没有修理！保安还查房卡入门时间，就是没有其他人在深夜进入我的房间。难道内贼不会用其他高明的方式进入吗？我的羽绒服也被这个内贼放在地上！我醒来时没有多想！近中午时我才发觉钱包少了现金和信用卡！还有，这家酒店的态度很差！没有同情心！我之前授权的2000元，我打了国际电话，银行说两天前我入酒店的2000元授权了，可是酒店的财务不领情，说中国银行没有授权。我又打了国际电话，我的银行说通过了！这家4星级的酒店不负责，认为不可能发生，我报案了，我下次再也不住这个，酒店[MASK]，1星不到的服务态度，很可耻！我还要把这个事件说给那些想定这个酒店的住客。酒店为何停电，摄像头坏得也太凑巧了来让大家知道这种内贼行为是要强力打击的。好了，不说了！！！千元丢了小事。酒店的处理态度我很反感！我强力告诉大家和提醒其他人不要到该酒店！', '宾馆在小街道上，不大好找，但还好北京热心同胞很多~宾馆设施跟介绍的差不多，房间很小，确实挺小，但加上低价位因素，还是无超所值的；环境不错，就在小胡同内，安静整洁，暖气好足-_-||。。。呵还有一大优势就是从宾馆出发，步行不到十分钟就可以到梅兰芳故居等等，京味小胡同，北海距离好近呢。总之，不错。推荐给节约消费的，酒店[MASK]，自助游朋友~比较划算，附近特色小吃很多~']\n",
      "['地理位置还不错，到哪里都比较方便，但是服务不象是豪生集团管理的，比较差。下午睡了一觉并洗了一个澡，本来想让酒店再来打扫，酒店差，一下，所以，打开了，请打扫的服务灯，可是到晚上回酒店，发现打扫得服务灯被关掉了，而房间还是没有打扫过。', '我这次是第5次住在长春的雁鸣湖大酒店。昨晚夜里停电。深夜我睡着了。我的钱包被内贼进入我的房间，偷了我近1000元和4张信用卡。。。我的证件和外币，数码相机等都在房间的保险箱里，原封不动。我打了好几个小时的长途电话来处理我的信用卡的冻结。我报案了，这个4星酒店的保安摄像探头竟然坏了，没有修理！保安还查房卡入门时间，就是没有其他人在深夜进入我的房间。难道内贼不会用其他高明的方式进入吗？我的羽绒服也被这个内贼放在地上！我醒来时没有多想！近中午时我才发觉钱包少了现金和信用卡！还有，这家酒店的态度很差！没有同情心！我之前授权的2000元，我打了国际电话，银行说两天前我入酒店的2000元授权了，可是酒店的财务不领情，说中国银行没有授权。我又打了国际电话，我的银行说通过了！这家4星级的酒店不负责，认为不可能发生，我报案了，我下次再也不住这个，酒店差，1星不到的服务态度，很可耻！我还要把这个事件说给那些想定这个酒店的住客。酒店为何停电，摄像头坏得也太凑巧了来让大家知道这种内贼行为是要强力打击的。好了，不说了！！！千元丢了小事。酒店的处理态度我很反感！我强力告诉大家和提醒其他人不要到该酒店！', '宾馆在小街道上，不大好找，但还好北京热心同胞很多~宾馆设施跟介绍的差不多，房间很小，确实挺小，但加上低价位因素，还是无超所值的；环境不错，就在小胡同内，安静整洁，暖气好足-_-||。。。呵还有一大优势就是从宾馆出发，步行不到十分钟就可以到梅兰芳故居等等，京味小胡同，北海距离好近呢。总之，不错。推荐给节约消费的，酒店好，自助游朋友~比较划算，附近特色小吃很多~']\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "text = []\n",
    "label = []\n",
    "punc = \"＂!＃＄％＆＇?（）()/＊＋，－／：；,.＜＝＞＠［＼］\\\"＾＿｀｛｜｝～｟｠｢｣､　、〃〈〉《》「」『』【】〔〕〖〗〘〙〚〛〜〝〞〟〰〾〿–—‘’‛“”„‟…‧﹏﹑﹔·！？｡。\"\n",
    "for index, row in tqdm(iterable=df_train.iterrows(), total=df_train.shape[0]):\n",
    "    sentence = row[\"review\"]\n",
    "    words = jieba.lcut(sentence)\n",
    "    for i in range(len(words)):\n",
    "        sentence_train = \"\".join(words[:i])+\"，酒店[MASK]，\"+\"\".join(words[i:])\n",
    "        sentence_test = \"\".join(words[:i])+\"，酒店\"+id_2_label[str(row[\"label\"])]+\"，\"+\"\".join(words[i:])\n",
    "        text.append(sentence_train)\n",
    "        label.append(sentence_test)\n",
    "text, label = shuffle(text, label)\n",
    "print(len(text))\n",
    "print(text[:3])\n",
    "print(label[:3])\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "def dataset_builder(x: List[str], y: List[str], tokenizer: BertTokenizer, max_len: int) -> Dataset:\n",
    "    data_dict = {'text': x, 'label_text': y}\n",
    "    result = Dataset.from_dict(data_dict)\n",
    "    def preprocess_function(examples):\n",
    "        text_token = tokenizer(examples['text'], padding=True,truncation=True, max_length=max_len)\n",
    "        text_token['labels'] = np.array(tokenizer(examples['label_text'], padding=True,truncation=True, max_length=max_len)[\"input_ids\"])\n",
    "        return text_token\n",
    "    result = result.map(preprocess_function, batched=True)\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "dd0d8ca81bd340aca1a76fb8626b6c19",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1 [00:00<?, ?ba/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "641cebd7fdcc4a3a9fb48591215bdf4d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/2 [00:00<?, ?ba/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "eval_dataset = dataset_builder(text[:130], label[:130], tokenizer, 512)\n",
    "train_dataset = dataset_builder(text[130:], label[130:], tokenizer, 512)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "PyTorch: setting up devices\n",
      "The default value for the training argument `--report_to` will change in v5 (from all installed integrations to none). In v5, you will need to use `--report_to all` to get the same behavior as now. You should start updating your code and make this info disappear :-).\n"
     ]
    }
   ],
   "source": [
    "args = TrainingArguments(\n",
    "    output_dir=\"/data/christmas.wang/project/classification_base_project/output\",\n",
    "    evaluation_strategy=\"steps\",\n",
    "    eval_steps=100,\n",
    "    per_device_train_batch_size=16,\n",
    "    per_device_eval_batch_size=16,\n",
    "    learning_rate=1e-5,\n",
    "    num_train_epochs=6,\n",
    "    seed=20,\n",
    "    load_best_model_at_end=True,\n",
    ")\n",
    "\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=args,\n",
    "    train_dataset=train_dataset,\n",
    "    eval_dataset=eval_dataset,\n",
    "    compute_metrics=compute_metrics,\n",
    "    callbacks=[EarlyStoppingCallback(early_stopping_patience=3)],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The following columns in the training set  don't have a corresponding argument in `BertForMaskedLM.forward` and have been ignored: label_text, text.\n",
      "***** Running training *****\n",
      "  Num examples = 1225\n",
      "  Num Epochs = 6\n",
      "  Instantaneous batch size per device = 16\n",
      "  Total train batch size (w. parallel, distributed & accumulation) = 16\n",
      "  Gradient Accumulation steps = 1\n",
      "  Total optimization steps = 462\n",
      "Automatic Weights & Biases logging enabled, to disable set os.environ[\"WANDB_DISABLED\"] = \"true\"\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='462' max='462' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [462/462 05:17, Epoch 6/6]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: left;\">\n",
       "      <th>Step</th>\n",
       "      <th>Training Loss</th>\n",
       "      <th>Validation Loss</th>\n",
       "      <th>Accuracy</th>\n",
       "      <th>F1</th>\n",
       "      <th>Precision</th>\n",
       "      <th>Recall</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>100</td>\n",
       "      <td>No log</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>1.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>200</td>\n",
       "      <td>No log</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>1.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>300</td>\n",
       "      <td>No log</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>1.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>400</td>\n",
       "      <td>No log</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>1.000000</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The following columns in the evaluation set  don't have a corresponding argument in `BertForMaskedLM.forward` and have been ignored: label_text, text.\n",
      "***** Running Evaluation *****\n",
      "  Num examples = 130\n",
      "  Batch size = 16\n",
      "Saving model checkpoint to /data/christmas.wang/project/classification_base_project/output/checkpoint-100\n",
      "Configuration saved in /data/christmas.wang/project/classification_base_project/output/checkpoint-100/config.json\n",
      "Model weights saved in /data/christmas.wang/project/classification_base_project/output/checkpoint-100/pytorch_model.bin\n",
      "The following columns in the evaluation set  don't have a corresponding argument in `BertForMaskedLM.forward` and have been ignored: label_text, text.\n",
      "***** Running Evaluation *****\n",
      "  Num examples = 130\n",
      "  Batch size = 16\n",
      "Saving model checkpoint to /data/christmas.wang/project/classification_base_project/output/checkpoint-200\n",
      "Configuration saved in /data/christmas.wang/project/classification_base_project/output/checkpoint-200/config.json\n",
      "Model weights saved in /data/christmas.wang/project/classification_base_project/output/checkpoint-200/pytorch_model.bin\n",
      "The following columns in the evaluation set  don't have a corresponding argument in `BertForMaskedLM.forward` and have been ignored: label_text, text.\n",
      "***** Running Evaluation *****\n",
      "  Num examples = 130\n",
      "  Batch size = 16\n",
      "Saving model checkpoint to /data/christmas.wang/project/classification_base_project/output/checkpoint-300\n",
      "Configuration saved in /data/christmas.wang/project/classification_base_project/output/checkpoint-300/config.json\n",
      "Model weights saved in /data/christmas.wang/project/classification_base_project/output/checkpoint-300/pytorch_model.bin\n",
      "The following columns in the evaluation set  don't have a corresponding argument in `BertForMaskedLM.forward` and have been ignored: label_text, text.\n",
      "***** Running Evaluation *****\n",
      "  Num examples = 130\n",
      "  Batch size = 16\n",
      "Saving model checkpoint to /data/christmas.wang/project/classification_base_project/output/checkpoint-400\n",
      "Configuration saved in /data/christmas.wang/project/classification_base_project/output/checkpoint-400/config.json\n",
      "Model weights saved in /data/christmas.wang/project/classification_base_project/output/checkpoint-400/pytorch_model.bin\n",
      "\n",
      "\n",
      "Training completed. Do not forget to share your model on huggingface.co/models =)\n",
      "\n",
      "\n",
      "Loading best model from /data/christmas.wang/project/classification_base_project/output/checkpoint-400 (score: 1.1244792830211736e-07).\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=462, training_loss=6.282679510839057e-05, metrics={'train_runtime': 321.0703, 'train_samples_per_second': 22.892, 'train_steps_per_second': 1.439, 'total_flos': 1839892910335200.0, 'train_loss': 6.282679510839057e-05, 'epoch': 6.0})"
      ]
     },
     "execution_count": 67,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████████████████████| 3999/3999 [00:48<00:00, 82.80it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'accuracy': 0.5356339084771192, 'f1': 0.5460767538499144, 'precision': 0.5341941654710665, 'recall': 0.5585}\n",
      "['多', '多', '多', '多', '豪', '居', '重', '多', '多', '全', '多', '微', '居', '重', '远', '电', '重', '多', '多', '微', '明', '惠', '[PAD]', '[PAD]', '多', '微', '多', '店', '光', '台', '全', '高', '远', '全', '超', '微', '多', '多', '多', '高', '尾', '[PAD]', '高', '多', '多', '全', '多', '全', '店', '店', '提', '部', '多', '角', '店', '超', '全', '大', '多', '店', '[PAD]', '多', '店', '远', '高', '店', '装', '店', '结', '全', '高', '多', '明', '多', '多', '近', '多', '全', '超', '多', '微', '多', '全', '多', '微', '多', '多', '店', '多', '台', '大', '微', '家', '店', '店', '明', '内', '店', '多', '尾', '有', '店', '多', '多', '多', '大', '多', '豪', '超', '店', '高', '店', '多', '全', '店', '高', '有', '多', '[PAD]', '微', '[PAD]', '多', '面', '重', '多', '正', '全', '大', '响', '多', '角', '瑞', '多', '有', '全', '全', '多', '粗', '多', '店', '多', '庄', '惠', '装', '重', '檔', '装', '多', '店', '多', '高', '全', '高', '多', '微', '店', '店', '多', '微', '全', '高', '有', '重', '等', '惠', '多', '多', '多', '装', '店', '高', '多', '台', '多', '店', '多', '多', '多', '全', '多', '位', '共', '店', '多', '位', '超', '重', '多', '多', '多', '店', '位', '底', '根', '微', '有', '高', '多', '多', '踏', '子', '尾', '多', '多', '远', '店', '多', '店', '多', '店', '角', '全', '远', '店', '重', '高', '多', '有', '多', '店', '全', '全', '店', '多', '尾', '光', '[PAD]', '多', '结', '多', '店', '高', '多', '多', '多', '多', '全', '尾', '全', '多', '点', '店', '亮', '饰', '多', '锁', '高', '多', '店', '多', '明', '全', '多', '多', '多', '多', '多', '多', '店', '[PAD]', '顶', '等', '店', '足', '多', '豪', '店', '[PAD]', '全', '多', '提', '店', '多', '多', '多', '[PAD]', '全', '早', '多', '全', '近', '多', '店', '店', '沈', '多', '远', '多', '多', '明', '全', '店', '分', '店', '多', '豪', '多', '示', '全', '惠', '丰', '有', '明', '大', '高', '多', '[PAD]', '装', '惠', '丰', '惠', '多', '角', '多', '全', '多', '光', '店', '沈', '多', '多', '多', '部', '多', '越', '全', '正', '多', '微', '多', '大', '林', '等', '尾', '响', '高', '多', '微', '店', '全', '重', '多', '多', '多', '大', '店', '豪', '多', '角', '多', '多', '多', '多', '惠', '明', '[PAD]', '重', '店', '台', '多', '微', '[PAD]', '锐', '店', '高', '林', '店', '光', '多', '豪', '拨', '多', '超', '多', '多', '越', '全', '微', '店', '华', '[PAD]', '多', '明', '台', '点', '有', '尾', '全', '多', '多', '底', '印', '重', '多', '多', '多', '家', '多', '沈', '全', '全', '多', '店', '全', '全', '多', '多', '多', '多', '多', '重', '多', '多', '位', '装', '位', '多', '多', '多', '高', '多', '店', '丰', '高', '豪', '多', '多', '店', '多', '高', '多', '高', '脚', '惠', '耳', '有', '微', '多', '响', '多', '有', '惠', '多', '多', '居', '高', '微', '多', '多', '全', '重', '明', '多', '多', '全', '多', '多', '全', '重', '[PAD]', '尾', '光', '多', '大', '多', '远', '大', '多', '尾', '微', '辉', '多', '多', '多', '店', '放', '全', '多', '尾', '高', '分', '多', '高', '全', '多', '尾', '远', '座', '全', '多', '多', '多', '多', '全', '尾', '多', '甸', '多', '多', '吉', '越', '多', '多', '店', '丰', '等', '远', '多', '大', '服', '多', '多', '超', '明', '尾', '明', '微', '微', '多', '早', '精', '多', '[PAD]', '店', '等', '店', '有', '店', '全', '多', '多', '多', '超', '多', '尾', '多', '全', '多', '重', '多', '锐', '林', '店', '多', '全', '多', '高', '等', '多', '锐', '全', '多', '装', '店', '超', '多', '超', '角', '明', '惠', '多', '重', '立', '多', '多', '越', '多', '高', '全', '籍', '多', '店', '多', '店', '高', '全', '提', '多', '店', '装', '多', '越', '多', '超', '多', '台', '角', '高', '装', '全', '店', '全', '店', '电', '多', '尾', '多', '店', '明', '多', '多', '装', '高', '店', '多', '店', '重', '印', '多', '多', '多', '远', '装', '多', '装', '踏', '重', '居', '多', '店', '多', '豪', '高', '多', '豪', '装', '有', '全', '响', '店', '多', '多', '多', '丰', '大', '尾', '光', '[PAD]', '长', '超', '多', '多', '多', '超', '尾', '店', '多', '全', '多', '多', '多', '多', '足', '高', '越', '高', '位', '有', '大', '尾', '遍', '等', '多', '多', '尾', '角', '重', '多', '越', '装', '微', '味', '高', '微', '多', '惠', '子', '多', '店', '全', '多', '重', '超', '远', '全', '店', '小', '多', '多', '頂', '微', '清', '多', '多', '超', '店', '高', '多', '多', '多', '店', '高', '籍', '重', '声', '远', '多', '尾', '亮', '多', '响', '多', '[PAD]', '甸', '底', '多', '惠', '多', '全', '多', '多', '高', '店', '全', '店', '底', '醒', '全', '全', '多', '惠', '全', '高', '尾', '超', '丰', '惠', '多', '多', '超', '多', '微', '有', '店', '多', '光', '全', '结', '多', '重', '店', '[UNK]', '尾', '多', '高', '多', '尾', '装', '店', '尾', '装', '多', '尾', '响', '多', '多', '高', '吉', '锐', '多', '多', '多', '店', '多', '多', '大', '华', '多', '多', '装', '底', '高', '多', '店', '豪', '高', '大', '惠', '大', '高', '豪', '多', '多', '豪', '微', '院', '微', '多', '多', '达', '多', '高', '多', '大', '多', '位', '店', '明', '尾', '微', '多', '多', '多', '店', '店', '有', '店', '端', '泰', '多', '明', '越', '全', '华', '多', '多', '装', '全', '多', '多', '多', '微', '多', '多', '不', '店', '超', '多', '全', '店', '多', '多', '多', '店', '多', '[UNK]', '尾', '店', '及', '尾', '多', '多', '有', '重', '明', '远', '[PAD]', '响', '全', '[PAD]', '多', '多', '有', '多', '店', '消', '微', '位', '多', '全', '微', '多', '店', '微', '店', '中', '全', '全', '沈', '多', '多', '飽', '超', '店', '多', '房', '多', '多', '高', '大', '大', '全', '装', '店', '多', '多', '甸', '高', '微', '豪', '多', '多', '大', '醒', '店', '多', '店', '多', '多', '越', '多', '多', '多', '微', '多', '店', '越', '集', '多', '大', '多', '装', '多', '店', '辉', '瑞', '提', '店', '惠', '多', '远', '多', '籍', '远', '微', '多', '露', '沙', '多', '丰', '足', '店', '微', '超', '全', '多', '店', '丰', '线', '全', '店', '店', '全', '家', '多', '多', '全', '店', '露', '多', '角', '惠', '店', '尾', '全', '客', '多', '多', '店', '頂', '多', '多', '多', '全', '多', '多', '尾', '多', '多', '店', '全', '多', '尾', '客', '庄', '惠', '位', '全', '店', '全', '多', '店', '店', '高', '多', '明', '店', '台', '明', '亮', '亮', '雅', '多', '多', '高', '台', '多', '集', '多', '惠', '多', '尾', '尾', '脚', '集', '多', '店', '多', '多', '多', '全', '超', '重', '全', '晚', '豪', '高', '多', '[PAD]', '店', '多', '微', '全', '明', '多', '大', '正', '多', '超', '多', '多', '末', '多', '林', '多', '店', '醒', '惠', '足', '高', '店', '微', '多', '多', '高', '多', '尾', '多', '多', '提', '多', '微', '尾', '多', '重', '多', '装', '微', '多', '磬', '多', '位', '全', '踏', '全', '多', '装', '多', '台', '多', '豪', '店', '越', '多', '装', '多', '越', '多', '明', '多', '多', '全', '多', '多', '明', '角', '多', '多', '明', '大', '店', '惠', '越', '重', '全', '多', '多', '分', '多', '全', '分', '多', '多', '越', '多', '全', '多', '多', '远', '豪', '境', '全', '角', '远', '店', '多', '底', '苏', '正', '居', '末', '多', '正', '多', '装', '电', '多', '光', '多', '全', '雅', '灿', '多', '多', '足', '惠', '店', '店', '多', '多', '等', '惠', '正', '多', '多', '底', '正', '多', '多', '足', '多', '多', '明', '灿', '多', '全', '多', '店', '多', '多', '全', '吉', '多', '全', '全', '越', '许', '多', '多', '多', '集', '等', '瑞', '多', '露', '超', '多', '远', '超', '装', '全', '微', '全', '店', '店', '1956', '多', '多', '多', '多', '多', '全', '多', '惠', '路', '多', '多', '多', '多', '客', '结', '全', '有', '惠', '高', '多', '装', '正', '多', '装', '多', '店', '多', '[PAD]', '底', '多', '多', '惠', '露', '多', '微', '位', '大', '有', '居', '有', '多', '多', '拨', '多', '全', '瑞', '高', '店', '高', '店', '大', '多', '标', '西', '家', '明', '位', '长', '多', '达', '多', '全', '苏', '多', '多', '亮', '多', '多', '多', '多', '多', '籍', '多', '大', '全', '店', '多', '子', '多', '全', '大', '店', '装', '高', '醒', '璜', '尾', '店', '多', '清', '多', '东', '端', '多', '有', '可', '多', '尾', '多', '多', '多', '微', '店', '多', '台', '多', '店', '微', '店', '重', '稍', '多', '多', '远', '全', '全', '微', '惠', '多', '店', '丰', '多', '多', '店', '多', '装', '全', '店', '价', '多', '有', '分', '多', '店', '多', '店', '一', '多', '多', '位', '多', '重', '多', '超', '醒', '光', '全', '居', '多', '宾', '全', '多', '大', '有', '大', '門', '多', '多', '全', '严', '超', '丰', '多', '店', '光', '多', '端', '多', '居', '早', '店', '多', '全', '惠', '多', '多', '藉', '店', '惠', '多', '多', '有', '光', '多', '全', '多', '超', '惠', '店', '座', '店', '全', '装', '有', '多', '多', '多', '有', '及', '多', '多', '多', '角', '多', '明', '有', '点', '店', '多', '多', '明', '多', '店', '多', '了', '分', '惠', '全', '多', '多', '多', '多', '多', '[PAD]', '解', '店', '全', '多', '多', '明', '多', '店', '尾', '全', '多', '微', '全', '全', '多', '多', '有', '大', '多', '多', '多', '瑞', '惠', '多', '多', '多', '反', '底', '店', '多', '清', '尾', '装', '正', '台', '香', '足', '多', '位', '阳', '装', '[PAD]', '全', '全', '多', '多', '多', '店', '价', '多', '台', '多', '吉', '多', '多', '多', '全', '全', '瑞', '多', '店', '多', '明', '全', '多', '可', '多', '多', '多', '多', '华', '店', '店', '高', '全', '多', '多', '远', '全', '明', '多', '多', '店', '高', '有', '越', '店', '尾', '惠', '大', '微', '微', '惠', '捷', '店', '惠', '多', '林', '莊', '尾', '苏', '音', '微', '豪', '店', '重', '微', '多', '重', '装', '明', '店', '阳', '多', '多', '超', '苏', '多', '多', '多', '全', '惠', '全', '店', '高', '多', '角', '装', '全', '店', '多', '有', '多', '多', '店', '端', '位', '店', '多', '多', '多', '多', '端', '多', '店', '惠', '惠', '惠', '多', '店', '尾', '全', '多', '明', '全', '多', '印', '多', '微', '分', '全', '足', '店', '足', '店', '锐', '嗨', '店', '丰', '尾', '微', '超', '店', '全', '微', '店', '位', '多', '多', '豪', '全', '店', '足', '店', '明', '多', '多', '多', '可', '全', '多', '全', '微', '店', '全', '多', '多', '厅', '全', '店', '遍', '装', '店', '台', '店', '大', '装', '店', '莊', '装', '高', '店', '装', '越', '多', '锐', '的', '店', '多', '瑞', '提', '超', '店', '多', '全', '店', '店', '豪', '台', '多', '全', '等', '多', '多', '多', '台', '惠', '重', '頂', '多', '远', '多', '多', '清', '多', '明', '多', '多', '超', '多', '亮', '光', '多', '多', '远', '高', '多', '高', '多', '微', '多', '吉', '富', '多', '多', '全', '多', '微', '多', '店', '家', '重', '全', '多', '装', '店', '全', '多', '多', '高', '多', '店', '遠', '尾', '店', '光', '多', '多', '多', '外', '多', '店', '多', '高', '尾', '店', '家', '多', '多', '出', '全', '家', '多', '多', '店', '全', '超', '多', '子', '多', '分', '青', '高', '多', '吉', '多', '高', '店', '多', '超', '豪', '丰', '全', '东', '多', '店', '澳', '明', '多', '全', '尾', '多', '多', '多', '光', '惠', '高', '多', '端', '多', '微', '店', '多', '多', '多', '多', '店', '超', '多', '亮', '全', '店', '端', '超', '多', '尾', '店', '多', '多', '外', '沈', '多', '重', '居', '超', '高', '高', '店']\n",
      "1788\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "pred = []\n",
    "true = []\n",
    "external_words = []\n",
    "df_test.dropna(how=\"any\", axis=0, inplace=True)\n",
    "for index, row in tqdm(iterable=df_test.iterrows(), total=df_test.shape[0]):\n",
    "    text = \"酒店[MASK]，\" + row[\"review\"]\n",
    "    tokenized_text = tokenizer.tokenize(text)\n",
    "    if len(tokenized_text) > 512:\n",
    "        tokenized_text = tokenized_text[:512]\n",
    "    indexed_tokens = tokenizer.convert_tokens_to_ids(tokenized_text)\n",
    "\n",
    "    # Create the segments tensors.\n",
    "    segments_ids = [0] * len(tokenized_text)\n",
    "    \n",
    "    # Convert inputs to PyTorch tensors\n",
    "    tokens_tensor = torch.tensor([indexed_tokens]).to('cuda')\n",
    "    segments_tensors = torch.tensor([segments_ids]).to('cuda')\n",
    "    \n",
    "    masked_index = tokenized_text.index('[MASK]')\n",
    "    \n",
    "    # Predict all tokens\n",
    "    with torch.no_grad():\n",
    "        predictions = model(tokens_tensor, segments_tensors)\n",
    "    predicted_index = torch.argmax(predictions[0][0][masked_index]).item()\n",
    "    predicted_token = tokenizer.convert_ids_to_tokens([predicted_index])[0]\n",
    "    # print(predicted_token+str(row[\"label\"]))\n",
    "    if predicted_token not in [\"差\", \"好\"]:\n",
    "        external_words.append(predicted_token)\n",
    "        predicted_token = \"差\"\n",
    "    y_pred = label_2_id[predicted_token]\n",
    "    pred.append(y_pred)\n",
    "    true.append(row[\"label\"])\n",
    "precision, recall, f1, _ = precision_recall_fscore_support(true, pred, average='binary')\n",
    "acc = accuracy_score(true, pred)\n",
    "print({'accuracy': acc, 'f1': f1, 'precision': precision, 'recall': recall})\n",
    "print(external_words)\n",
    "print(len(external_words))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [],
   "source": [
    "good_words = [\"美\", \"舒\", \"服\", \"豪\", \"华\", \"丽\", \"亮\", \"错\", \"大\", \"宜\", \"明\"]\n",
    "def get_label(words: List[str]) -> int:\n",
    "    for key, val in label_2_id.items():\n",
    "        if key in words:\n",
    "            return val\n",
    "    for word in words:\n",
    "        if word in good_words:\n",
    "            return 0\n",
    "    return 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████████████████████| 3999/3999 [00:48<00:00, 82.02it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'accuracy': 0.5253813453363341, 'f1': 0.6699130434782609, 'precision': 0.5136, 'recall': 0.963}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "pred = []\n",
    "true = []\n",
    "\n",
    "for index, row in tqdm(iterable=df_test.iterrows(), total=df_test.shape[0]):\n",
    "    text = \"酒店[MASK]，\" + row[\"review\"]\n",
    "    tokenized_text = tokenizer.tokenize(text)\n",
    "    if len(tokenized_text) > 512:\n",
    "        tokenized_text = tokenized_text[:512]\n",
    "    indexed_tokens = tokenizer.convert_tokens_to_ids(tokenized_text)\n",
    "\n",
    "    # Create the segments tensors.\n",
    "    segments_ids = [0] * len(tokenized_text)\n",
    "    \n",
    "    # Convert inputs to PyTorch tensors\n",
    "    tokens_tensor = torch.tensor([indexed_tokens]).to('cuda')\n",
    "    segments_tensors = torch.tensor([segments_ids]).to('cuda')\n",
    "    \n",
    "    masked_index = tokenized_text.index('[MASK]')\n",
    "    \n",
    "    # Predict all tokens\n",
    "    with torch.no_grad():\n",
    "        predictions = model(tokens_tensor, segments_tensors)\n",
    "    top_k = torch.topk(predictions[0][0][masked_index].flatten(), 5).indices.tolist()\n",
    "    words = []\n",
    "    for word in top_k:\n",
    "        predicted_token = tokenizer.convert_ids_to_tokens([word])[0]\n",
    "        words.append(predicted_token)\n",
    "    pred.append(get_label(words))\n",
    "    true.append(row[\"label\"])\n",
    "precision, recall, f1, _ = precision_recall_fscore_support(true, pred, average='binary')\n",
    "acc = accuracy_score(true, pred)\n",
    "print({'accuracy': acc, 'f1': f1, 'precision': precision, 'recall': recall})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_mask(data: List[Dict])->None:\n",
    "    for ele in tqdm(data):\n",
    "        sentence = ele[\"text\"].strip()\n",
    "        new_label = \"酒店\"+id_2_label[str(ele[\"label\"])]+\"，\"+sentence\n",
    "        new_sentence = \"酒店[MASK]，\"+sentence\n",
    "        ele[\"new_text\"] = new_sentence\n",
    "        ele[\"label_text\"] = new_label\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>label</th>\n",
       "      <th>review</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>1</td>\n",
       "      <td>价格比比较不错的酒店。这次免费升级了，感谢前台服务员。房子还好，地毯是新的，比上次的好些。早...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>1</td>\n",
       "      <td>宾馆在小街道上，不大好找，但还好北京热心同胞很多~宾馆设施跟介绍的差不多，房间很小，确实挺小...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>距离川沙公路较近,但是公交指示不对,如果是\"蔡陆线\"的话,会非常麻烦.建议用别的路线.房间较...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5326</th>\n",
       "      <td>0</td>\n",
       "      <td>我这次是第5次住在长春的雁鸣湖大酒店。昨晚夜里停电。深夜我睡着了。我的钱包被内贼进入我的房间...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5323</th>\n",
       "      <td>0</td>\n",
       "      <td>服务态度极其差，前台接待好象没有受过培训，连基本的礼貌都不懂，竟然同时接待几个客人；大堂副理...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>1</td>\n",
       "      <td>CBD中心,周围没什么店铺,说5星有点勉强.不知道为什么卫生间没有电吹风</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>1</td>\n",
       "      <td>不错，在同等档次酒店中应该是值得推荐的！</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>1</td>\n",
       "      <td>早餐太差，无论去多少人，那边也不加食品的。酒店应该重视一下这个问题了。房间本身很好。</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>1</td>\n",
       "      <td>1。酒店比较新，装潢和设施还不错，只是房间有些油漆味。2。早餐还可以，只是品种不是很多。3。...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5327</th>\n",
       "      <td>0</td>\n",
       "      <td>前台checkin花了20分钟，checkout25分钟，这是服务态度和没有做到位。信用卡刷...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5330</th>\n",
       "      <td>0</td>\n",
       "      <td>我3.6预定好的180的标间，当我到的时候竟然说有会议房间满了，我订的房间没有了，太不讲信誉...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>1</td>\n",
       "      <td>入住丽晶，感觉很好。因为是新酒店，的确有淡淡的油漆味，房间内较新。房间大小合适，卫生间设备齐...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5322</th>\n",
       "      <td>0</td>\n",
       "      <td>标准间太差房间还不如3星的而且设施非常陈旧.建议酒店把老的标准间从新改善.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5324</th>\n",
       "      <td>0</td>\n",
       "      <td>地理位置还不错，到哪里都比较方便，但是服务不象是豪生集团管理的，比较差。下午睡了一觉并洗了一...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5329</th>\n",
       "      <td>0</td>\n",
       "      <td>酒店服务态度极差，设施很差，建议还是不要到那儿去。</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5328</th>\n",
       "      <td>0</td>\n",
       "      <td>有或者很少房!梯部不吸,但是有一些吸者仍然有服!我是不抽的人,成二手的受害者!(中13人口中...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1</td>\n",
       "      <td>商务大床房，房间很大，床有2M宽，整体感觉经济实惠不错!</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5331</th>\n",
       "      <td>0</td>\n",
       "      <td>房间的环境非常差,而且房间还不隔音，住的不舒服。</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>1</td>\n",
       "      <td>总的来说，这样的酒店配这样的价格还算可以，希望他赶快装修，给我的客人留些好的印象</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5325</th>\n",
       "      <td>0</td>\n",
       "      <td>1。我住的是靠马路的标准间。房间内设施简陋，并且的房间玻璃窗户外还有一层幕墙玻璃，而且不能打...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "      label                                             review\n",
       "6         1  价格比比较不错的酒店。这次免费升级了，感谢前台服务员。房子还好，地毯是新的，比上次的好些。早...\n",
       "3         1  宾馆在小街道上，不大好找，但还好北京热心同胞很多~宾馆设施跟介绍的差不多，房间很小，确实挺小...\n",
       "0         1  距离川沙公路较近,但是公交指示不对,如果是\"蔡陆线\"的话,会非常麻烦.建议用别的路线.房间较...\n",
       "5326      0  我这次是第5次住在长春的雁鸣湖大酒店。昨晚夜里停电。深夜我睡着了。我的钱包被内贼进入我的房间...\n",
       "5323      0  服务态度极其差，前台接待好象没有受过培训，连基本的礼貌都不懂，竟然同时接待几个客人；大堂副理...\n",
       "4         1               CBD中心,周围没什么店铺,说5星有点勉强.不知道为什么卫生间没有电吹风\n",
       "7         1                               不错，在同等档次酒店中应该是值得推荐的！\n",
       "2         1         早餐太差，无论去多少人，那边也不加食品的。酒店应该重视一下这个问题了。房间本身很好。\n",
       "9         1  1。酒店比较新，装潢和设施还不错，只是房间有些油漆味。2。早餐还可以，只是品种不是很多。3。...\n",
       "5327      0  前台checkin花了20分钟，checkout25分钟，这是服务态度和没有做到位。信用卡刷...\n",
       "5330      0  我3.6预定好的180的标间，当我到的时候竟然说有会议房间满了，我订的房间没有了，太不讲信誉...\n",
       "8         1  入住丽晶，感觉很好。因为是新酒店，的确有淡淡的油漆味，房间内较新。房间大小合适，卫生间设备齐...\n",
       "5322      0              标准间太差房间还不如3星的而且设施非常陈旧.建议酒店把老的标准间从新改善.\n",
       "5324      0  地理位置还不错，到哪里都比较方便，但是服务不象是豪生集团管理的，比较差。下午睡了一觉并洗了一...\n",
       "5329      0                          酒店服务态度极差，设施很差，建议还是不要到那儿去。\n",
       "5328      0  有或者很少房!梯部不吸,但是有一些吸者仍然有服!我是不抽的人,成二手的受害者!(中13人口中...\n",
       "1         1                       商务大床房，房间很大，床有2M宽，整体感觉经济实惠不错!\n",
       "5331      0                           房间的环境非常差,而且房间还不隔音，住的不舒服。\n",
       "5         1           总的来说，这样的酒店配这样的价格还算可以，希望他赶快装修，给我的客人留些好的印象\n",
       "5325      0  1。我住的是靠马路的标准间。房间内设施简陋，并且的房间玻璃窗户外还有一层幕墙玻璃，而且不能打..."
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "origin_data_path = \"/data/tmp/nlp-data/open_source_data/classification/ChnSentiCorp_htl_all.csv\"\n",
    "df_origin_data = pd.read_csv(origin_data_path, encoding='utf-8')\n",
    "df_good = df_origin_data[df_origin_data[\"label\"] == 0].head(10)\n",
    "df_bad = df_origin_data[df_origin_data[\"label\"] == 1].head(10)\n",
    "\n",
    "df_few = shuffle(pd.concat([df_good,df_bad]))\n",
    "df_few"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████| 4765/4765 [00:00<00:00, 271222.70it/s]\n",
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████| 1500/1500 [00:00<00:00, 388265.61it/s]\n",
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████| 1500/1500 [00:00<00:00, 437849.26it/s]\n"
     ]
    }
   ],
   "source": [
    "add_mask(train_json)\n",
    "add_mask(eval_json)\n",
    "add_mask(test_json)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Didn't find file /data/tmp/christmas.wang/chinese_wwm_ext_pytorch/added_tokens.json. We won't load it.\n",
      "Didn't find file /data/tmp/christmas.wang/chinese_wwm_ext_pytorch/special_tokens_map.json. We won't load it.\n",
      "Didn't find file /data/tmp/christmas.wang/chinese_wwm_ext_pytorch/tokenizer_config.json. We won't load it.\n",
      "Didn't find file /data/tmp/christmas.wang/chinese_wwm_ext_pytorch/tokenizer.json. We won't load it.\n",
      "loading file /data/tmp/christmas.wang/chinese_wwm_ext_pytorch/vocab.txt\n",
      "loading file None\n",
      "loading file None\n",
      "loading file None\n",
      "loading file None\n",
      "loading configuration file /data/tmp/christmas.wang/chinese_wwm_ext_pytorch/config.json\n",
      "loading configuration file /data/tmp/christmas.wang/chinese_wwm_ext_pytorch/config.json\n",
      "Model config BertConfig {\n",
      "  \"attention_probs_dropout_prob\": 0.1,\n",
      "  \"classifier_dropout\": null,\n",
      "  \"directionality\": \"bidi\",\n",
      "  \"gradient_checkpointing\": false,\n",
      "  \"hidden_act\": \"gelu\",\n",
      "  \"hidden_dropout_prob\": 0.1,\n",
      "  \"hidden_size\": 768,\n",
      "  \"initializer_range\": 0.02,\n",
      "  \"intermediate_size\": 3072,\n",
      "  \"layer_norm_eps\": 1e-12,\n",
      "  \"max_position_embeddings\": 512,\n",
      "  \"model_type\": \"bert\",\n",
      "  \"num_attention_heads\": 12,\n",
      "  \"num_hidden_layers\": 12,\n",
      "  \"pad_token_id\": 0,\n",
      "  \"pooler_fc_size\": 768,\n",
      "  \"pooler_num_attention_heads\": 12,\n",
      "  \"pooler_num_fc_layers\": 3,\n",
      "  \"pooler_size_per_head\": 128,\n",
      "  \"pooler_type\": \"first_token_transform\",\n",
      "  \"position_embedding_type\": \"absolute\",\n",
      "  \"transformers_version\": \"4.10.2\",\n",
      "  \"type_vocab_size\": 2,\n",
      "  \"use_cache\": true,\n",
      "  \"vocab_size\": 21128\n",
      "}\n",
      "\n",
      "loading weights file /data/tmp/christmas.wang/chinese_wwm_ext_pytorch/pytorch_model.bin\n",
      "Some weights of the model checkpoint at /data/tmp/christmas.wang/chinese_wwm_ext_pytorch were not used when initializing BertForMaskedLM: ['cls.seq_relationship.weight', 'cls.seq_relationship.bias']\n",
      "- This IS expected if you are initializing BertForMaskedLM 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 BertForMaskedLM from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "All the weights of BertForMaskedLM were initialized from the model checkpoint at /data/tmp/christmas.wang/chinese_wwm_ext_pytorch.\n",
      "If your task is similar to the task the model of the checkpoint was trained on, you can already use BertForMaskedLM for predictions without further training.\n"
     ]
    }
   ],
   "source": [
    "tokenizer = BertTokenizer.from_pretrained(bert_name_path)\n",
    "model = BertForMaskedLM.from_pretrained(bert_name_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def dataset_builder(data: List[Dict], tokenizer: BertTokenizer, max_len: int) -> Dataset:\n",
    "    x = []\n",
    "    y = []\n",
    "    for ele in tqdm(data):\n",
    "        x.append(ele[\"new_text\"])\n",
    "        y.append(ele[\"label_text\"])\n",
    "    data_dict = {'text': x, 'label_text': y}\n",
    "    result = Dataset.from_dict(data_dict)\n",
    "    def preprocess_function(examples):\n",
    "        text_token = tokenizer(examples['text'], padding=True,truncation=True, max_length=max_len)\n",
    "        text_token['labels'] = np.array(tokenizer(examples['label_text'], padding=True,truncation=True, max_length=max_len)[\"input_ids\"])\n",
    "        return text_token\n",
    "    result = result.map(preprocess_function, batched=True)\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████| 4765/4765 [00:00<00:00, 638566.64it/s]\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d4f78c77044d4a4e910a9db16345ade0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?ba/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████████████████████████████████████████████████████████| 1500/1500 [00:00<00:00, 1064364.07it/s]\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fade11d237454973985667b54dfaa6ec",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/2 [00:00<?, ?ba/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████| 1500/1500 [00:00<00:00, 587108.62it/s]\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "30167a3e895c49e38fd4f2bb84a9a991",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/2 [00:00<?, ?ba/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "Dataset({\n",
       "    features: ['attention_mask', 'input_ids', 'label_text', 'labels', 'text', 'token_type_ids'],\n",
       "    num_rows: 1500\n",
       "})"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_dataset = dataset_builder(train_json, tokenizer, 64)\n",
    "eval_dataset = dataset_builder(eval_json, tokenizer, 64)\n",
    "test_dataset = dataset_builder(test_json, tokenizer, 64)\n",
    "test_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_metrics(pred):\n",
    "    labels = pred.label_ids[:, 3]\n",
    "    preds = pred.predictions[:, 3].argmax(-1)\n",
    "    precision, recall, f1, _ = precision_recall_fscore_support(labels, preds, average='weighted')\n",
    "    acc = accuracy_score(labels, preds)\n",
    "    return {'accuracy': acc, 'f1': f1, 'precision': precision, 'recall': recall}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "PyTorch: setting up devices\n",
      "The default value for the training argument `--report_to` will change in v5 (from all installed integrations to none). In v5, you will need to use `--report_to all` to get the same behavior as now. You should start updating your code and make this info disappear :-).\n"
     ]
    }
   ],
   "source": [
    "args = TrainingArguments(\n",
    "    output_dir=\"/data/christmas.wang/project/classification_base_project/output\",\n",
    "    evaluation_strategy=\"steps\",\n",
    "    eval_steps=500,\n",
    "    per_device_train_batch_size=16,\n",
    "    per_device_eval_batch_size=16,\n",
    "    learning_rate=1e-5,\n",
    "    num_train_epochs=6,\n",
    "    seed=20,\n",
    "    load_best_model_at_end=True,\n",
    ")\n",
    "\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=args,\n",
    "    train_dataset=test_dataset,\n",
    "    eval_dataset=eval_dataset,\n",
    "    compute_metrics=compute_metrics,\n",
    "    callbacks=[EarlyStoppingCallback(early_stopping_patience=3)],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import wandb\n",
    "# wandb.init()\n",
    "\n",
    "# trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████| 4765/4765 [00:48<00:00, 98.59it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'accuracy': 0.3246589716684155, 'f1': 0.024257125530624625, 'precision': 0.8695652173913043, 'recall': 0.012300123001230012}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "pred = []\n",
    "true = []\n",
    "external_words = []\n",
    "\n",
    "\n",
    "counter = 0\n",
    "for ele in tqdm(train_json):\n",
    "    if counter > 2:\n",
    "        break\n",
    "    tokenized_text = tokenizer.tokenize(ele['new_text'])\n",
    "    if len(tokenized_text) > 512:\n",
    "        tokenized_text = tokenized_text[:512]\n",
    "    indexed_tokens = tokenizer.convert_tokens_to_ids(tokenized_text)\n",
    "\n",
    "    # Create the segments tensors.\n",
    "    segments_ids = [0] * len(tokenized_text)\n",
    "    \n",
    "    # Convert inputs to PyTorch tensors\n",
    "    tokens_tensor = torch.tensor([indexed_tokens]).to('cuda')\n",
    "    segments_tensors = torch.tensor([segments_ids]).to('cuda')\n",
    "    \n",
    "    masked_index = tokenized_text.index('[MASK]')\n",
    "    \n",
    "    # Predict all tokens\n",
    "    with torch.no_grad():\n",
    "        predictions = model(tokens_tensor, segments_tensors)\n",
    "    predicted_index = torch.argmax(predictions[0][0][masked_index]).item()\n",
    "    predicted_token = tokenizer.convert_ids_to_tokens([predicted_index])[0]\n",
    "    # print(predicted_token+str(row[\"label\"]))\n",
    "    if predicted_token not in [\"差\", \"好\"]:\n",
    "        external_words.append(predicted_token)\n",
    "        predicted_token = \"差\"\n",
    "    y_pred = label_2_id[predicted_token]\n",
    "    pred.append(y_pred)\n",
    "    true.append(ele[\"label\"])\n",
    "precision, recall, f1, _ = precision_recall_fscore_support(true, pred, average='binary')\n",
    "acc = accuracy_score(true, pred)\n",
    "print({'accuracy': acc, 'f1': f1, 'precision': precision, 'recall': recall})\n",
    "# print(external_words)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "good_words = [\"美\", \"舒\", \"服\", \"豪\", \"华\", \"丽\", \"亮\"]\n",
    "def get_label(words: List[str]) -> int:\n",
    "    for key, val in label_2_id.items():\n",
    "        if key in words:\n",
    "            return val\n",
    "    for word in words:\n",
    "        if word in good_words:\n",
    "            return 0\n",
    "    return 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████| 4765/4765 [00:48<00:00, 98.47it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'accuracy': 0.683105981112277, 'f1': 0.8113914564076942, 'precision': 0.6832141354648716, 'recall': 0.998769987699877}\n",
      "[]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "pred = []\n",
    "true = []\n",
    "external_words = []\n",
    "\n",
    "\n",
    "counter = 0\n",
    "for ele in tqdm(train_json):\n",
    "    if counter > 2:\n",
    "        break\n",
    "    tokenized_text = tokenizer.tokenize(ele['new_text'])\n",
    "    if len(tokenized_text) > 512:\n",
    "        tokenized_text = tokenized_text[:512]\n",
    "    indexed_tokens = tokenizer.convert_tokens_to_ids(tokenized_text)\n",
    "\n",
    "    # Create the segments tensors.\n",
    "    segments_ids = [0] * len(tokenized_text)\n",
    "    \n",
    "    # Convert inputs to PyTorch tensors\n",
    "    tokens_tensor = torch.tensor([indexed_tokens]).to('cuda')\n",
    "    segments_tensors = torch.tensor([segments_ids]).to('cuda')\n",
    "    \n",
    "    masked_index = tokenized_text.index('[MASK]')\n",
    "    \n",
    "    # Predict all tokens\n",
    "    with torch.no_grad():\n",
    "        predictions = model(tokens_tensor, segments_tensors)\n",
    "    top_k = torch.topk(predictions[0][0][masked_index].flatten(), 5).indices.tolist()\n",
    "    words = []\n",
    "    for word in top_k:\n",
    "        predicted_token = tokenizer.convert_ids_to_tokens([word])[0]\n",
    "        words.append(predicted_token)\n",
    "    pred.append(get_label(words))\n",
    "    true.append(ele[\"label\"])\n",
    "precision, recall, f1, _ = precision_recall_fscore_support(true, pred, average='binary')\n",
    "acc = accuracy_score(true, pred)\n",
    "print({'accuracy': acc, 'f1': f1, 'precision': precision, 'recall': recall})\n",
    "print(external_words)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.10"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}