re/make_aug_data.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "from easydict import EasyDict\n",
    "import logging\n",
    "import os\n",
    "from typing import List, Tuple\n",
    "from utils import InputExample, InputFeatures\n",
    "import json\n",
    "import pandas as pd\n",
    "import random"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "args = EasyDict({\n",
    "    \"batch_size\": 32,\n",
    "    \"data_dir\" : \"./data\",\n",
    "    \"model_dir\": \"./model\",\n",
    "    \"model_tarname\":\"klue-re.tar.gz\",\n",
    "    \"output_dir\":os.environ.get(\"SM_OUTPUT_DATA_DIR\", \"/output\"),\n",
    "    \"max_seq_length\":512,\n",
    "    \"relation_filename\" : \"relation_list.json\",\n",
    "    \"train_filename\" : \"klue-re-v1.1_train.json\",\n",
    "    \"valid_filename\" : \"klue-re-v1.1_dev.json\",\n",
    "    \"num_workers\" : 4\n",
    "})\n",
    "# 릴레이션 데이터 위치\n",
    "relation_class_file_path = os.path.join(args.data_dir, args.relation_filename)\n",
    "# train 데이터 위치\n",
    "train_file_path = os.path.join(args.data_dir, args.train_filename)\n",
    "aug_entity_swap_file_path = os.path.join(args.data_dir, 'train_aug_entity_swap.json')\n",
    "aug_aeda_file_path = os.path.join(args.data_dir, 'train_aug_aeda.json')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(relation_class_file_path, \"r\", encoding=\"utf-8\") as f:\n",
    "    relation_class = json.load(f)[\"relations\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_df = pd.read_json(train_file_path,orient='recode')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "def aeda(row,p,punctuations):\n",
    "    sentence = row['sentence']\n",
    "    sub_entity = row['subject_entity']\n",
    "    obj_entity = row['object_entity']\n",
    "\n",
    "    sub_start = sub_entity['start_idx']\n",
    "    sub_end = sub_entity['end_idx']\n",
    "\n",
    "    obj_start = obj_entity['start_idx']\n",
    "    obj_end = obj_entity['end_idx']\n",
    "    new_sentence =''\n",
    "    sub_add = 0\n",
    "    obj_add = 0\n",
    "    for i in range(len(sentence)):\n",
    "        if sentence[i] != ' ':\n",
    "            new_sentence += sentence[i]\n",
    "        else:\n",
    "            prob = random.random()\n",
    "            if prob < p:\n",
    "\n",
    "                if not ((sub_start <= i <= sub_end) or (obj_start <= i <= obj_end)):\n",
    "                    punc_idx = random.randint(0,len(punctuations) - 1)\n",
    "                    add_punc = ' '+punctuations[punc_idx]+' '\n",
    "                    new_sentence += add_punc\n",
    "                    \n",
    "                    if sub_end <= obj_start:\n",
    "                        if i <= sub_start:\n",
    "                            sub_add += 2\n",
    "                            obj_add += 2\n",
    "\n",
    "                        elif sub_end <= i <= obj_start:\n",
    "                            obj_add += 2\n",
    "                    elif obj_end <= sub_start:\n",
    "                        if i <= obj_start:\n",
    "                            sub_add += 2\n",
    "                            obj_add += 2\n",
    "                        elif obj_end <= i <= sub_start:\n",
    "                            sub_add += 2\n",
    "                else:\n",
    "                    new_sentence += ' '\n",
    "            else:\n",
    "                new_sentence += ' '\n",
    "        \n",
    "    return new_sentence, sub_add, obj_add"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "def entity_swap(row, new_label):\n",
    "    sentence = row['sentence']\n",
    "    sub_entity = row['subject_entity']\n",
    "    obj_entity = row['object_entity']\n",
    "\n",
    "    n_data = {}\n",
    "    n_data[\"guid\"] = row[\"guid\"]\n",
    "    n_data[\"sentence\"] = sentence\n",
    "    n_data[\"subject_entity\"] = obj_entity\n",
    "    n_data[\"object_entity\"] = sub_entity\n",
    "    n_data['label'] = new_label\n",
    "    n_data['source'] = row['source']\n",
    "\n",
    "    return n_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "punctuations=[\".\", \",\", \"!\", \"?\", \";\", \":\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "aeda_list = [\n",
    "        #\"no_relation\",\n",
    "        \"org:dissolved\",\n",
    "        \"org:founded\",\n",
    "        \"org:place_of_headquarters\",\n",
    "        \"org:alternate_names\",\n",
    "        \"org:member_of\",\n",
    "        \"org:members\",\n",
    "        \"org:political/religious_affiliation\",\n",
    "        \"org:product\",\n",
    "        \"org:founded_by\",\n",
    "        \"org:top_members/employees\",\n",
    "        \"org:number_of_employees/members\",\n",
    "        \"per:date_of_birth\",\n",
    "        \"per:date_of_death\",\n",
    "        \"per:place_of_birth\",\n",
    "        \"per:place_of_death\",\n",
    "        \"per:place_of_residence\",\n",
    "        \"per:origin\",\n",
    "        \"per:employee_of\",\n",
    "        \"per:schools_attended\",\n",
    "        \"per:alternate_names\",\n",
    "        \"per:parents\",\n",
    "        \"per:children\",\n",
    "        \"per:siblings\",\n",
    "        \"per:spouse\",\n",
    "        \"per:other_family\",\n",
    "        \"per:colleagues\",\n",
    "        \"per:product\",\n",
    "        \"per:religion\",\n",
    "        \"per:title\"\n",
    "    ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "rel_data = train_df[train_df['label'] == \"org:member_of\"].reset_index()\n",
    "aug_data1 = []\n",
    "for i in range(len(rel_data)):\n",
    "    aug_data1.append(entity_swap(rel_data.iloc[i], \"org:members\"))\n",
    "\n",
    "aug_data2 = []\n",
    "rel_data = train_df[train_df['label'] == \"org:members\"].reset_index()\n",
    "for i in range(len(rel_data)):\n",
    "    aug_data2.append(entity_swap(rel_data.iloc[i], \"org:member_of\"))\n",
    "\n",
    "aug_data3 = []\n",
    "rel_data = train_df[train_df['label'] == \"per:other_family\"].reset_index()\n",
    "for i in range(len(rel_data)):\n",
    "    aug_data3.append(entity_swap(rel_data.iloc[i], \"per:other_family\"))\n",
    "\n",
    "aug_data4 = []\n",
    "rel_data = train_df[train_df['label'] == \"per:colleagues\"].reset_index()\n",
    "for i in range(len(rel_data)):\n",
    "    aug_data4.append(entity_swap(rel_data.iloc[i], \"per:colleagues\"))\n",
    "\n",
    "aug_data5 = []\n",
    "rel_data = train_df[train_df['label'] == \"per:parents\"].reset_index()\n",
    "for i in range(len(rel_data)):\n",
    "    aug_data5.append(entity_swap(rel_data.iloc[i], \"per:children\"))\n",
    "\n",
    "aug_data6 = []\n",
    "rel_data = train_df[train_df['label'] == \"per:children\"].reset_index()\n",
    "for i in range(len(rel_data)):\n",
    "    aug_data6.append(entity_swap(rel_data.iloc[i], \"per:parents\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "aug_data = []\n",
    "aug_data.extend(aug_data1)\n",
    "aug_data.extend(aug_data2)\n",
    "aug_data.extend(aug_data3)\n",
    "aug_data.extend(aug_data4)\n",
    "aug_data.extend(aug_data5)\n",
    "aug_data.extend(aug_data6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3834"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(aug_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(aug_entity_swap_file_path, \"w\") as f:\n",
    "    json.dump(aug_data, f,ensure_ascii=False,indent=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "aeda_data = []\n",
    "\n",
    "for relation in aeda_list:\n",
    "    rel_data = train_df[train_df['label'] == relation].reset_index()\n",
    "    for i in range(len(rel_data)):\n",
    "        #print(rel_data.iloc[i])\n",
    "        ns, sa, oa = aeda(rel_data.iloc[i],0.7,punctuations)\n",
    "        sub_entity = rel_data.iloc[i][\"subject_entity\"]\n",
    "        obj_entity = rel_data.iloc[i][\"object_entity\"]\n",
    "        n_data = {}\n",
    "        n_data[\"guid\"] = rel_data.iloc[i][\"guid\"]\n",
    "        n_data[\"sentence\"] = ns\n",
    "        n_data[\"subject_entity\"] = {\n",
    "            \"word\" : sub_entity['word'], ###\n",
    "            \"start_idx\" : sub_entity[\"start_idx\"] + sa,\n",
    "            \"end_idx\" : sub_entity[\"end_idx\"] + sa,\n",
    "            \"type\" : sub_entity[\"type\"]\n",
    "        }\n",
    "        n_data[\"object_entity\"] = {\n",
    "            \"word\" : obj_entity['word'], ###\n",
    "            \"start_idx\" : obj_entity[\"start_idx\"] + oa,\n",
    "            \"end_idx\" : obj_entity[\"end_idx\"] + oa,\n",
    "            \"type\" : obj_entity[\"type\"]\n",
    "        }\n",
    "        n_data['label'] = rel_data.iloc[i][\"label\"]\n",
    "        n_data['source'] = rel_data.iloc[i]['source']\n",
    "        aeda_data.append(n_data)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(aug_aeda_file_path, \"w\") as f:\n",
    "    json.dump(aeda_data, f,ensure_ascii=False,indent=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(len(aeda_data)):\n",
    "    sen = aeda_data[i]['sentence']\n",
    "    sub_e = aeda_data[i]['subject_entity']\n",
    "    obj_e = aeda_data[i]['object_entity']\n",
    "    #print(sen)\n",
    "\n",
    "    obj_word = obj_e['word']\n",
    "    o_s = obj_e['start_idx']\n",
    "    o_e = obj_e['end_idx']\n",
    "\n",
    "    sub_word = sub_e['word']\n",
    "    s_s = sub_e['start_idx']\n",
    "    s_e = sub_e['end_idx']\n",
    "\n",
    "    if sub_word != sen[s_s : s_e + 1]:\n",
    "        print(i)\n",
    "        break\n",
    "    if obj_word != sen[o_s : o_e + 1]:\n",
    "        print(i)\n",
    "        break\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "22936"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(aeda_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "d4d1e4263499bec80672ea0156c357c1ee493ec2b1c70f0acce89fc37c4a6abe"
  },
  "kernelspec": {
   "display_name": "Python 3.8.5 64-bit ('base': conda)",
   "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.12"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}