3.CHIP-CDN-PredictTextsim.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import pickle\n",
    "import pandas as pd\n",
    "\n",
    "from ark_nlp.model.tm.bert import Bert\n",
    "from ark_nlp.model.tm.bert import BertConfig\n",
    "from ark_nlp.model.tm.bert import Dataset\n",
    "from ark_nlp.model.tm.bert import Task\n",
    "from ark_nlp.model.tm.bert import get_default_model_optimizer\n",
    "from ark_nlp.model.tm.bert import Tokenizer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 一、数据读入与处理"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### 1. 召回模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import copy\n",
    "import logging\n",
    "import numpy as np\n",
    "\n",
    "from six import iteritems\n",
    "\n",
    "\n",
    "logger = logging.getLogger(__name__)\n",
    "\n",
    "\n",
    "class BM25(object):\n",
    "    \"\"\"\n",
    "    BM25模型\n",
    "\n",
    "    Args:\n",
    "        corpus (:obj:`list`):\n",
    "            检索的语料\n",
    "        k1 (:obj:`float`, optional, defaults to 1.5):\n",
    "            取正值的调优参数，用于文档中的词项频率进行缩放控制\n",
    "        b (:obj:`float`, optional, defaults to 0.75):\n",
    "            0到1之间的参数，决定文档长度的缩放程度，b=1表示基于文档长度对词项权重进行完全的缩放，b=0表示归一化时不考虑文档长度因素\n",
    "        epsilon (:obj:`float`, optional, defaults to 0.25):\n",
    "            idf的下限值\n",
    "        tokenizer (:obj:`object`, optional, defaults to None):\n",
    "            分词器，用于对文档进行分词操作，默认为None，按字颗粒对文档进行分词\n",
    "        is_retain_docs (:obj:`bool`, optional, defaults to False):\n",
    "            是否保持原始文档\n",
    "\n",
    "    Reference:\n",
    "        [1] https://github.com/RaRe-Technologies/gensim/blob/3.8.3/gensim/summarization/bm25.py\n",
    "    \"\"\"  # noqa: ignore flake8\"\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        corpus,\n",
    "        k1=1.5,\n",
    "        b=0.75,\n",
    "        epsilon=0.25,\n",
    "        tokenizer=None,\n",
    "        is_retain_docs=False\n",
    "    ):\n",
    "        self.k1 = k1\n",
    "        self.b = b\n",
    "        self.epsilon = epsilon\n",
    "\n",
    "        self.docs = None\n",
    "        self.corpus_size = 0\n",
    "        self.avgdl = 0\n",
    "        self.doc_freqs = []\n",
    "        self.idf = {}\n",
    "        self.doc_len = []\n",
    "\n",
    "        if is_retain_docs:\n",
    "            self.docs = copy.deepcopy(corpus)\n",
    "\n",
    "        if tokenizer:\n",
    "            corpus = [self.tokenizer.tokenize(document) for document in corpus]\n",
    "        else:\n",
    "            corpus = [list(document) for document in corpus]\n",
    "\n",
    "        self._initialize(corpus)\n",
    "\n",
    "    def _initialize(self, corpus):\n",
    "        \"\"\"Calculates frequencies of terms in documents and in corpus. Also computes inverse document frequencies.\"\"\"\n",
    "        nd = {}  # word -> number of documents with word\n",
    "        num_doc = 0\n",
    "        for document in corpus:                        \n",
    "            self.corpus_size += 1\n",
    "            self.doc_len.append(len(document))\n",
    "            num_doc += len(document)\n",
    "\n",
    "            frequencies = {}\n",
    "            for word in document:\n",
    "                if word not in frequencies:\n",
    "                    frequencies[word] = 0\n",
    "                frequencies[word] += 1\n",
    "            self.doc_freqs.append(frequencies)\n",
    "\n",
    "            for word, freq in iteritems(frequencies):\n",
    "                if word not in nd:\n",
    "                    nd[word] = 0\n",
    "                nd[word] += 1\n",
    "\n",
    "        self.avgdl = float(num_doc) / self.corpus_size\n",
    "\n",
    "        idf_sum = 0\n",
    "        negative_idfs = []\n",
    "        for word, freq in iteritems(nd):\n",
    "            idf = math.log(self.corpus_size - freq + 0.5) - math.log(freq + 0.5)\n",
    "            self.idf[word] = idf\n",
    "            idf_sum += idf\n",
    "            if idf < 0:\n",
    "                negative_idfs.append(word)\n",
    "        self.average_idf = float(idf_sum) / len(self.idf)\n",
    "\n",
    "        if self.average_idf < 0:\n",
    "            logger.warning(\n",
    "                'Average inverse document frequency is less than zero. Your corpus of {} documents'\n",
    "                ' is either too small or it does not originate from natural text. BM25 may produce'\n",
    "                ' unintuitive results.'.format(self.corpus_size)\n",
    "            )\n",
    "\n",
    "        eps = self.epsilon * self.average_idf\n",
    "        for word in negative_idfs:\n",
    "            self.idf[word] = eps\n",
    "\n",
    "    def get_score(self, query, index):\n",
    "        score = 0.0\n",
    "        doc_freqs = self.doc_freqs[index]\n",
    "        numerator_constant = self.k1 + 1\n",
    "        denominator_constant = self.k1 * (1 - self.b + self.b * self.doc_len[index] / self.avgdl)\n",
    "        for word in query:\n",
    "            if word in doc_freqs:\n",
    "                df = self.doc_freqs[index][word]\n",
    "                idf = self.idf[word]\n",
    "                score += (idf * df * numerator_constant) / (df + denominator_constant)\n",
    "        return score\n",
    "\n",
    "    def get_scores(self, query):\n",
    "        scores = [self.get_score(query, index) for index in range(self.corpus_size)]\n",
    "        return scores\n",
    "\n",
    "    def recall(self, query, topk=5):\n",
    "        scores = self.get_scores(query)\n",
    "        indexs = np.argsort(scores)[::-1][:topk]\n",
    "\n",
    "        if self.docs is None:\n",
    "            return [[i, scores[i]] for i in indexs]\n",
    "        else:\n",
    "            return [[self.docs[i], scores[i]] for i in indexs]\n",
    "\n",
    "bm25_model = pickle.load(open('../checkpoint/recall/bm25_model.pkl', 'rb'))\n",
    "map_dict = pickle.load(open('../checkpoint/recall/map_dict.pkl', 'rb'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 2. 数据生成"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data_df = pd.read_json('../data/source_datasets/CHIP-CDN/CHIP-CDN_train.json')\n",
    "dev_data_df = pd.read_json('../data/source_datasets/CHIP-CDN/CHIP-CDN_dev.json')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pair_dataset = []\n",
    "for _raw_word, _normalized_result in zip(train_data_df['text'], train_data_df['normalized_result']):\n",
    "    normalized_words = set(_normalized_result.split('##'))\n",
    "    search_result_ = set()\n",
    "    train_pair_dataset = []\n",
    "    for _index, _search_word in enumerate(\n",
    "        [_result for _results in bm25_model.recall(_raw_word, topk=1000) for _result in map_dict[_results[0]]]):\n",
    "\n",
    "        if _search_word in normalized_words:\n",
    "            continue\n",
    "        elif _search_word in search_result_:\n",
    "            continue\n",
    "        else:\n",
    "            train_pair_dataset.append([_raw_word, _search_word, '0'])\n",
    "            \n",
    "        search_result_.add(_search_word)\n",
    "            \n",
    "        if len(train_pair_dataset) == 20:\n",
    "            pair_dataset.extend(train_pair_dataset)\n",
    "            break\n",
    "                    \n",
    "    for _st_word in normalized_words:\n",
    "        for _ in range(10):\n",
    "            pair_dataset.append([_raw_word, _st_word, '1'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pair_dev_dataset = []\n",
    "for _raw_word, _normalized_result in zip(train_data_df['text'], train_data_df['normalized_result']):\n",
    "    normalized_words = set(_normalized_result.split('##'))\n",
    "    search_result_ = set()\n",
    "    dev_pair_dataset = []\n",
    "    for _index, _search_word in enumerate(\n",
    "        [_result for _results in bm25_model.recall(_raw_word, topk=1000) for _result in map_dict[_results[0]]]):\n",
    "\n",
    "        if _search_word in normalized_words:\n",
    "            continue\n",
    "        elif _search_word in search_result_:\n",
    "            continue\n",
    "        else:\n",
    "            dev_pair_dataset.append([_raw_word, _search_word, '0'])\n",
    "            \n",
    "        search_result_.add(_search_word)\n",
    "            \n",
    "        if len(dev_pair_dataset) == 1:\n",
    "            pair_dev_dataset.extend(dev_pair_dataset)\n",
    "            break\n",
    "                    \n",
    "    for _st_word in normalized_words:\n",
    "        pair_dev_dataset.append([_raw_word, _st_word, '1'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data_df = pd.DataFrame(pair_dataset, columns=['text_a', 'text_b', 'label'])\n",
    "dev_data_df = pd.DataFrame(pair_dev_dataset, columns=['text_a', 'text_b', 'label'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tm_train_dataset = Dataset(train_data_df)\n",
    "tm_dev_dataset = Dataset(dev_data_df)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 2. 词典创建"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = Tokenizer(vocab='nghuyong/ernie-1.0', max_seq_len=50)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 3. 生成分词器"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 4. ID化"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tm_train_dataset.convert_to_ids(tokenizer)\n",
    "tm_dev_dataset.convert_to_ids(tokenizer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>\n",
    "\n",
    "### 二、模型构建"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 1. 模型参数设置"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config = BertConfig.from_pretrained('nghuyong/ernie-1.0',\n",
    "                                    num_labels=len(tm_train_dataset.cat2id))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 2. 模型创建"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dl_module = Bert.from_pretrained('nghuyong/ernie-1.0', \n",
    "                                 config=config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>\n",
    "\n",
    "### 三、任务构建"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 1. 任务参数和必要部件设定"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 设置运行次数\n",
    "num_epoches = 2\n",
    "batch_size = 32"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "param_optimizer = list(dl_module.named_parameters())\n",
    "param_optimizer = [n for n in param_optimizer if 'pooler' not in n[0]]\n",
    "no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']\n",
    "optimizer_grouped_parameters = [\n",
    "    {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)],\n",
    "     'weight_decay': 0.01},\n",
    "    {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}\n",
    "]      "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 2. 任务创建"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Task(dl_module, 'adamw', 'lsce', cuda_device=0, ema_decay=0.995)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 3. 训练"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.fit(tm_train_dataset, \n",
    "          tm_dev_dataset,\n",
    "          lr=3e-5,\n",
    "          epochs=num_epoches, \n",
    "          batch_size=batch_size,\n",
    "          params=optimizer_grouped_parameters\n",
    "         )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.ema.store(model.module.parameters())\n",
    "model.ema.copy_to(model.module.parameters())  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>\n",
    "\n",
    "### 四、模型验证与保存"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 1. 模型验证"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ark_nlp.factory.predictor import TMPredictor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tm_predictor_instance = TMPredictor(model.module, tokenizer, tm_train_dataset.cat2id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tm_predictor_instance.predict_one_sample(['胸部皮肤破裂伤', '胸部开放性损伤'], return_proba=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 2. 模型保存"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(model.module.state_dict(), '../checkpoint/textsim/module.pth')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('../checkpoint/textsim/cat2id.pkl', \"wb\") as f:\n",
    "    pickle.dump(tm_train_dataset.cat2id, f)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}