CMeIE.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4c72901",
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "import os\n",
    "import json\n",
    "import jieba\n",
    "import torch\n",
    "import pickle\n",
    "import codecs\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import pandas as pd\n",
    "\n",
    "from ark_nlp.model.re.casrel_bert import CasRelBert\n",
    "from ark_nlp.model.re.casrel_bert import CasRelBertConfig\n",
    "from ark_nlp.model.re.casrel_bert import Dataset\n",
    "from ark_nlp.model.re.casrel_bert import Task\n",
    "from ark_nlp.model.re.casrel_bert import get_default_model_optimizer\n",
    "from ark_nlp.model.re.casrel_bert import Tokenizer\n",
    "from ark_nlp.factory.loss_function import CasRelLoss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69854942",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 目录地址\n",
    "\n",
    "train_data_path = '../data/source_datasets/CMeIE/CMeIE_train.json'\n",
    "dev_data_path = '../data/source_datasets/CMeIE/CMeIE_dev.json'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc598e6a",
   "metadata": {},
   "source": [
    "#### 1. 数据读入"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77edd06b",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data_list = []\n",
    "\n",
    "with codecs.open(train_data_path, mode='r', encoding='utf8') as f:\n",
    "    lines = f.readlines()\n",
    "    for index_, line_ in enumerate(lines):\n",
    "        record_ = {}\n",
    "        line_ = json.loads(line_.strip())\n",
    "        record_['text'] = line_['text']\n",
    "        record_['label'] = []\n",
    "        for triple_ in line_['spo_list']:\n",
    "            record_['label'].append([\n",
    "                triple_['subject'],\n",
    "                record_['text'].index(triple_['subject']),\n",
    "                record_['text'].index(triple_['subject'])+ len(triple_['subject']) - 1,\n",
    "                triple_['predicate'] + '@' + triple_['object_type']['@value'],\n",
    "                triple_['object']['@value'],\n",
    "                record_['text'].index(triple_['object']['@value']),\n",
    "                record_['text'].index(triple_['object']['@value']) + len(triple_['object']['@value']) - 1,\n",
    "            ])\n",
    "        train_data_list.append(record_)\n",
    "\n",
    "train_df = pd.DataFrame(train_data_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a443cad",
   "metadata": {},
   "outputs": [],
   "source": [
    "dev_data_list = []\n",
    "counter = 0\n",
    "with codecs.open(dev_data_path, mode='r', encoding='utf8') as f:\n",
    "    lines = f.readlines()\n",
    "    for index_, line_ in enumerate(lines):\n",
    "        record_ = {}\n",
    "        line_ = json.loads(line_.strip())\n",
    "        record_['text'] = line_['text']\n",
    "        record_['label'] = []\n",
    "        for triple_ in line_['spo_list']:\n",
    "            record_['label'].append([\n",
    "                triple_['subject'],\n",
    "                record_['text'].index(triple_['subject']),\n",
    "                record_['text'].index(triple_['subject'])+ len(triple_['subject']) - 1,\n",
    "                triple_['predicate'] + '@' + triple_['object_type']['@value'],\n",
    "                triple_['object']['@value'],\n",
    "                record_['text'].index(triple_['object']['@value']),\n",
    "                record_['text'].index(triple_['object']['@value']) + len(triple_['object']['@value']) - 1,\n",
    "            ])\n",
    "            counter += 1\n",
    "        dev_data_list.append(record_)\n",
    "        \n",
    "dev_df = pd.DataFrame(dev_data_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24d51c3c",
   "metadata": {},
   "outputs": [],
   "source": [
    "re_train_dataset = Dataset(train_df)\n",
    "re_dev_dataset = Dataset(dev_df,\n",
    "                         categories = re_train_dataset.categories,\n",
    "                         is_train=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eec01479",
   "metadata": {},
   "source": [
    "#### 2. 词典创建和生成分词器"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41bfb854",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = Tokenizer(\n",
    "    vocab='nghuyong/ernie-1.0',\n",
    "    max_seq_len=128\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67ecaa82",
   "metadata": {},
   "source": [
    "#### 3. ID化"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29d01f62",
   "metadata": {},
   "outputs": [],
   "source": [
    "re_train_dataset.convert_to_ids(tokenizer)\n",
    "re_dev_dataset.convert_to_ids(tokenizer)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8918cb18",
   "metadata": {},
   "source": [
    "<br>\n",
    "\n",
    "### 二、模型构建"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "71ef3699",
   "metadata": {},
   "source": [
    "#### 1. 模型参数设置"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30173f8f",
   "metadata": {},
   "outputs": [],
   "source": [
    "bert_config = CasRelBertConfig.from_pretrained('nghuyong/ernie-1.0',\n",
    "                                               num_labels=len(re_train_dataset.cat2id))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af4442d6",
   "metadata": {},
   "source": [
    "#### 2. 模型创建"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8882050a",
   "metadata": {},
   "outputs": [],
   "source": [
    "dl_module = CasRelBert.from_pretrained('nghuyong/ernie-1.0',\n",
    "                                       config=bert_config)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5187f461",
   "metadata": {},
   "source": [
    "<br>\n",
    "\n",
    "### 三、任务构建"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "589f7157",
   "metadata": {},
   "source": [
    "#### 1. 任务参数和必要部件设定"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e6c7add",
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = get_default_model_optimizer(dl_module) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aeb86665",
   "metadata": {},
   "source": [
    "#### 2. 任务创建"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e822c7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Task(dl_module, optimizer, CasRelLoss(), cuda_device=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5931e90",
   "metadata": {},
   "source": [
    "#### 3. 训练"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47ab7d32",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.fit(\n",
    "    re_train_dataset,\n",
    "    re_dev_dataset,\n",
    "    lr=4e-5,\n",
    "    epochs=20,\n",
    "    batch_size=16\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a13ba861",
   "metadata": {},
   "source": [
    "<br>\n",
    "\n",
    "### 四、模型预测"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d19146a-7b68-48ad-b6dc-af0ea83db410",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "from ark_nlp.model.re.casrel_bert import Predictor\n",
    "\n",
    "casrel_re_predictor_instance = Predictor(model.module, tokenizer, re_train_dataset.cat2id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80d6c230-6fd3-42ac-b343-0faa09acd0ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_data_path = '../data/source_datasets/CMeIE/CMeIE_test.json'\n",
    "schemas_data_path = '../data/source_datasets/CMeIE/53_schemas.json'\n",
    "output_data_path = '../data/output_datasets/CMeIE_test.json'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1f2d01c-455b-46a9-8acb-25b10a15f22f",
   "metadata": {},
   "outputs": [],
   "source": [
    "result = []\n",
    "\n",
    "with open(test_data_path, 'r', encoding='utf-8') as f:\n",
    "    lines = f.readlines()\n",
    "    for line_ in lines:\n",
    "        result.append(casrel_re_predictor_instance.predict_one_sample(eval(line_)['text']))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df0072e8-bc34-45db-8d5c-84e4dd24cb35",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_subject_type = []\n",
    "all_predicate = []\n",
    "all_shcemas = []\n",
    "predicate2subject = {}\n",
    "with open(schemas_data_path, 'r', encoding='utf-8') as fs:\n",
    "    for jsonstr in fs.readlines():\n",
    "        jsonstr = json.loads(jsonstr)\n",
    "        # all_shcemas.append(jsonstr)\n",
    "        \n",
    "        predicate2subject[jsonstr['predicate']+'@'+jsonstr['object_type']] = jsonstr['subject_type']\n",
    "        \n",
    "    fs.close()\n",
    "    \n",
    "all_predicate = set(all_predicate)\n",
    "with open(output_data_path, 'w', encoding='utf-8') as fw:\n",
    "    with open(test_data_path, 'r', encoding='utf-8') as f:\n",
    "        lines = f.readlines()\n",
    "        for index_, jsonstr in tqdm(enumerate(lines)):\n",
    "            line = json.loads(jsonstr)\n",
    "            results_len = []\n",
    "            sentence = line['text']\n",
    "            dict_list = result[index_]\n",
    "            new = []\n",
    "            for list_ in dict_list:\n",
    "                for predicate_ in predicate2subject:\n",
    "                    if list_[1] == predicate_:\n",
    "                        if list_[-1] != '' and list_[-1] != '[UNK]':\n",
    "                            result_dict = {\n",
    "                                'predicate': predicate_.split('@')[0],\n",
    "                                \"subject\": list_[0],\n",
    "                                'subject_type': predicate2subject[predicate_],\n",
    "                                \"object\": {\"@value\": list_[-1]},\n",
    "                                'object_type': {\"@value\":predicate_.split('@')[-1]}\n",
    "                                }\n",
    "                        else:\n",
    "                            continue\n",
    "                        if result_dict not in new:\n",
    "                            new.append(result_dict)\n",
    "            if sum([item.count('。') for item in sentence]) >= 2:\n",
    "                for item in new:\n",
    "                    item['Combined'] = True\n",
    "            else:\n",
    "                for item in new:\n",
    "                    item['Combined'] = False\n",
    "\n",
    "            if len(new) == 0:\n",
    "                new = [{\n",
    "                    \"Combined\": '',\n",
    "                    \"predicate\": '',\n",
    "                    \"subject\": '',\n",
    "                    \"subject_type\": '',\n",
    "                    \"object\": {\"@value\": \"\"},\n",
    "                    \"object_type\": {\"@value\": \"\"},\n",
    "                }]\n",
    "                pred_dict = {\n",
    "                    \"text\": ''.join(sentence),\n",
    "                    \"spo_list\": new,\n",
    "                }\n",
    "            else:\n",
    "\n",
    "                pred_dict = {\n",
    "                    \"text\": ''.join(sentence),\n",
    "                    \"spo_list\": new,\n",
    "                }\n",
    "            fw.write(json.dumps(pred_dict, ensure_ascii=False) + '\\n')\n",
    "f.close()\n",
    "fw.close()"
   ]
  }
 ],
 "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": 5
}