legal-nlp/06.2.Relation_Extraction_Training.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "i8W51t04BN6B"
      },
      "source": [
        "![JohnSnowLabs](https://nlp.johnsnowlabs.com/assets/images/logo.png)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "A5g1XTOJif3u"
      },
      "source": [
        "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/JohnSnowLabs/spark-nlp-workshop/blob/master/legal-nlp/06.2.Relation_Extraction_Training.ipynb)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Starting the spark session"
      ],
      "metadata": {
        "id": "_0p9WrsexKS7"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "pycharm": {
          "is_executing": true
        },
        "id": "_914itZsj51v"
      },
      "outputs": [],
      "source": [
        "! pip install -q johnsnowlabs"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "pycharm": {
          "is_executing": true
        },
        "id": "fY0lcShkj51w"
      },
      "outputs": [],
      "source": [
        "from johnsnowlabs import *\n",
        "\n",
        "# nlp.install(force_browser=True)"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from google.colab import files\n",
        "print('Please Upload your John Snow Labs License using the button below')\n",
        "license_keys = files.upload()"
      ],
      "metadata": {
        "id": "i57QV3-_P2sQ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "nlp.install()"
      ],
      "metadata": {
        "id": "OfmmPqknP4rR"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wRXTnNl3j51w"
      },
      "outputs": [],
      "source": [
        "from johnsnowlabs import *\n",
        "spark = nlp.start()"
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "7ZOX73gkOmOB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3R2CM28aC-wB"
      },
      "source": [
        "# Relation Extraction training using TensorFlow 2.x and BERT"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import tensorflow.compat.v1 as tf\n"
      ],
      "metadata": {
        "id": "z7nKiI8A1umb"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!wget https://storage.googleapis.com/bert_models/2018_10_18/cased_L-12_H-768_A-12.zip\n",
        "\n"
      ],
      "metadata": {
        "id": "LRwrLUoizE1c"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!mkdir /content/trained\n",
        "!mkdir /content/models\n",
        "!mkdir /content/models/bert_base"
      ],
      "metadata": {
        "id": "giaebks6BQt8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!unzip /content/cased_L-12_H-768_A-12.zip -d /content/models/bert_base/"
      ],
      "metadata": {
        "id": "Q6vgU1Oo7esz"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## RE DL Training Notebook for Tensorflow 2\n",
        "\n",
        "An adaptation of the original Google Bert TF 1 repo "
      ],
      "metadata": {
        "id": "AUonFKDTUruU"
      }
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "p9a7QwiuC-wC"
      },
      "source": [
        "# 2. Download BERT code implementation and BERT weights\n",
        "In this section we will download official BERT code and the Bert pretrained weights we will use to finetune and create our RE model."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ASRUv_ZAouZM"
      },
      "source": [
        "## 2.1. Downloading BERT code"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VYv-kob8Nuvb"
      },
      "outputs": [],
      "source": [
        "#Bert source location\n",
        "BERT_SRC = \"./bert_t\"\n",
        "#Bert repo to download source from\n",
        "BERT_REPO = \"https://github.com/google-research/bert\""
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!wget  -q https://raw.githubusercontent.com/JohnSnowLabs/spark-nlp-workshop/master/legal-nlp/data/tf2_contrib.py\n",
        "!wget  -q https://raw.githubusercontent.com/JohnSnowLabs/spark-nlp-workshop/master/legal-nlp/data/relations.csv"
      ],
      "metadata": {
        "id": "m7siGZIdfRfP"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nc6rZCruNuvc",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "054f3248-397c-4ae5-dfd4-1f0345302135"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Cloning into './bert_t'...\n",
            "remote: Enumerating objects: 340, done.\u001b[K\n",
            "remote: Total 340 (delta 0), reused 0 (delta 0), pack-reused 340\u001b[K\n",
            "Receiving objects: 100% (340/340), 328.27 KiB | 16.41 MiB/s, done.\n",
            "Resolving deltas: 100% (182/182), done.\n"
          ]
        }
      ],
      "source": [
        "!test -d $BERT_SRC || git clone $BERT_REPO $BERT_SRC\n",
        "!sed -i 's/import tensorflow as tf/import tensorflow.compat.v1 as tf/g' $BERT_SRC/optimization.py\n",
        "!sed -i 's/import tensorflow as tf/import tensorflow.compat.v1 as tf/g' $BERT_SRC/run_classifier.py\n",
        "!sed -i 's/import tensorflow as tf/import tensorflow.compat.v1 as tf/g' $BERT_SRC/tokenization.py\n",
        "!sed -i 's/import tensorflow as tf/import tensorflow.compat.v1 as tf\\nfrom tf2_contrib import tf_contrib_layer_norm\\n/g' $BERT_SRC/modeling.py\n",
        "!sed -i 's/tf.contrib.layers.layer_norm/tf_contrib_layer_norm/g' $BERT_SRC/modeling.py\n",
        "!cp tf2_contrib.py $BERT_SRC/"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "plRnHzH7C-wE"
      },
      "source": [
        "## 2.2.Add BERT to System Path\n",
        "Bert code will look for several modules in the system path, so we need to set that they can be also find in `bert` folder"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EMu6cjHZNuvd"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if not BERT_SRC in sys.path:\n",
        "    sys.path += [BERT_SRC]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4v68yr0oC-wJ"
      },
      "source": [
        "## Imports"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Y_-NbGKJNuvd"
      },
      "outputs": [],
      "source": [
        "import tensorflow.compat.v1 as tf\n",
        "import pandas as pd\n",
        "import pickle\n",
        "import sys\n",
        "import re\n",
        "import json\n",
        "import numpy as np\n",
        "import modeling\n",
        "import optimization\n",
        "import tokenization\n",
        "import run_classifier\n",
        "import shutil\n",
        "import os\n",
        "import pprint\n",
        "from IPython.display import clear_output\n",
        "from scipy.spatial.distance import cosine, euclidean\n",
        "from functools import reduce\n",
        "import scipy.stats as stats\n",
        "import pyspark.sql.functions as F\n",
        "\n",
        "pp = pprint.PrettyPrinter(indent=4)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xqdZ1zP6GANF"
      },
      "source": [
        "## Hyperparam configuration\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UJb7DF_mGhD1"
      },
      "source": [
        "### Base BERT"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "code_folding": [
          76
        ],
        "id": "UIysdzmhNuve"
      },
      "outputs": [],
      "source": [
        " \n",
        "class BaseBertI2B2Config:\n",
        "    #maximum sequence length, can be up to 512 for standard Bertmodels\n",
        "    #larger values require more GPU memory\n",
        "    #A GTX1080 Ti with 11 GB of memory can do no more than batch size 16 with max_seq_len 128\n",
        "    MAX_SEQ_LENGTH = 128\n",
        "\n",
        "    #location of pretrained Bert model\n",
        "    BERT_MODEL_PATH = \"/content/models/bert_base/cased_L-12_H-768_A-12/\"\n",
        "    #location of Bert chekpoint used for initializing the model\n",
        "    BERT_MODEL_CHECKPOINT_PATH = \"{}/bert_model.ckpt\".format(BERT_MODEL_PATH)\n",
        "    #location of Bert configuration file\n",
        "    BERT_MODEL_CONFIG_PATH = \"{}/bert_config.json\".format(BERT_MODEL_PATH)\n",
        "    #location of Bert vocabulary file\n",
        "    BERT_VOCAB_PATH = \"{}/vocab.txt\".format(BERT_MODEL_PATH)\n",
        "\n",
        "    #Location for storing trained models  (in checkpoing format)\n",
        "    CHKPOINT_PATH = \"./trained\"\n",
        "   \n",
        "    #Location to export trained models to (in saved_model format)\n",
        "    EXPORT_PATH = \"/content/models/basebert_re\"\n",
        "\n",
        "    #Initial LR, real LR depends on warm-up and training progress\n",
        "    LEARNING_RATE = 2e-5\n",
        "    #Number of training epochs (how many time to iterate through the training set)\n",
        "    NUM_TRAIN_EPOCHS = 5\n",
        "    #Proportion of training steps(i.e. number of batches) used for warming up (adaptive LR in the begging)\n",
        "    WARMUP_PROPORTION = 0.1\n",
        "    #Training batch size\n",
        "    BATCH_SIZE = 16\n",
        "    #Batch size during testing/valdiation\n",
        "    V_BATCH_SIZE = 100\n",
        "\n",
        "    #Sentence column name\n",
        "    SENTENCE_COLUMN = \"text\"\n",
        "    #Relation label column name\n",
        "    REL_LABEL_COLUMN = \"rel\"\n",
        "    #Relation argument binding colum name - used if (some of the) relations are not symmetric\n",
        "    #0 - symmetric relation, argument order doesn't matter\n",
        "    #1 - rel(ARG1, ARG2), where ARGS1 is the entity which first appears in the text\n",
        "    #2 - rel(ARG2, ARG1)\n",
        "    #if None, then ignore argument order(i.e. treat all relations as symmetric)\n",
        "    REL_ARG_BINDING_COLUMN = None\n",
        "\n",
        "    #Entities positions in the dataset\n",
        "    ENTITY1_BEGIN_COLUMN = \"firstCharEnt1\"\n",
        "    ENTITY1_END_COLUMN = \"lastCharEnt1\"\n",
        "    ENTITY2_BEGIN_COLUMN = \"firstCharEnt2\"\n",
        "    ENTITY2_END_COLUMN = \"lastCharEnt2\"\n",
        "\n",
        "\n",
        "    ENTITY1_START_TAG = \"e1b\"\n",
        "    ENTITY1_END_TAG = \"e1e\"\n",
        "    ENTITY2_START_TAG = \"e2b\"\n",
        "    ENTITY2_END_TAG = \"e2e\"\n",
        "\n",
        "    ENTITY1_START_TAG_ID = 10\n",
        "    ENTITY1_END_TAG_ID = 11\n",
        "    ENTITY2_START_TAG_ID = 12\n",
        "    ENTITY2_END_TAG_ID = 13\n",
        "\n",
        "    #stadard padding id value for Bert models\n",
        "    PAD_ID = 0\n",
        "\n",
        "    #proportion of training examples\n",
        "    TRAIN_SET_PROB = 0.8\n",
        "    \n",
        "    #Not used at the moment  \n",
        "    NUM_HIDDEN_UNITS = 0\n",
        "    DROPOUT_RATE = 0\n",
        "    REPLACE_ARG_PROB = 0    \n",
        "    \n",
        "    USE_ENTITY_POSITIONS = True\n",
        "    USE_CLS_POSITION = True"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#BertREConfig is used by the code in this notebook, update for each model you train\n",
        "BertREConfig = BaseBertI2B2Config"
      ],
      "metadata": {
        "id": "1JEZJncmNL9v"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oEt6gSySHOO0"
      },
      "source": [
        "## Data collection\n",
        "Set of functions to get input data from pandas or Spark dataframes"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cda7eoHxHkfG"
      },
      "source": [
        "### Reading RE data from a pandas dataset"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "code_folding": [
          22,
          45
        ],
        "id": "BOWk2yBXNuvf"
      },
      "outputs": [],
      "source": [
        "def collect_data_from_pandas_dataset(dataset):\n",
        "    \n",
        "    rel_labels = sorted(dataset[BertREConfig.REL_LABEL_COLUMN].unique())\n",
        "    \n",
        "    def process_row(row):\n",
        "\n",
        "        row[\"sentence\"] = annotate_sentence(\n",
        "            row[BertREConfig.SENTENCE_COLUMN], \n",
        "            row[BertREConfig.ENTITY1_BEGIN_COLUMN],\n",
        "            row[BertREConfig.ENTITY1_END_COLUMN],\n",
        "            row[BertREConfig.ENTITY2_BEGIN_COLUMN],\n",
        "            row[BertREConfig.ENTITY2_END_COLUMN]\n",
        "        )\n",
        "        row[\"rel_label_id\"] = rel_labels.index(row[BertREConfig.REL_LABEL_COLUMN])\n",
        "        row[\"rel_arg_binding\"] = row[BertREConfig.REL_ARG_BINDING_COLUMN] if BertREConfig.REL_ARG_BINDING_COLUMN else 0\n",
        "\n",
        "        return row\n",
        "    \n",
        "    \n",
        "    dataset = dataset.apply(process_row, axis=1)\n",
        "    \n",
        "    return dataset.sentence, dataset.rel_label_id, dataset.rel_arg_binding, rel_labels\n",
        "    \n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oa6ZlI3RHoXp"
      },
      "source": [
        "### Reading RE data from a spark dataset"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "def collect_data_from_spark_dataset(dataset):\n",
        "    \n",
        "    rel_labels = sorted([row[0] for row in dataset.select(BertREConfig.REL_LABEL_COLUMN).distinct().collect()])\n",
        "    \n",
        "    def process_row(row):\n",
        "        sentence = annotate_sentence(\n",
        "            row[BertREConfig.SENTENCE_COLUMN], \n",
        "            int(row[BertREConfig.ENTITY1_BEGIN_COLUMN]),\n",
        "            int(row[BertREConfig.ENTITY1_END_COLUMN]),\n",
        "            int(row[BertREConfig.ENTITY2_BEGIN_COLUMN]),\n",
        "            int(row[BertREConfig.ENTITY2_END_COLUMN])\n",
        "        )       \n",
        "        rel_label_id = rel_labels.index(row[BertREConfig.REL_LABEL_COLUMN])\n",
        "        \n",
        "        rel_arg_binding = (\n",
        "            row[BertREConfig.REL_ARG_BINDING_COLUMN] if BertREConfig.REL_ARG_BINDING_COLUMN else 0)\n",
        "        \n",
        "        return (sentence, rel_label_id, rel_arg_binding)\n",
        "    \n",
        "    sentences, rel_label_ids, rel_arg_bindings = tuple(\n",
        "        map(list, zip(*dataset.rdd.map(process_row).collect())))\n",
        "    \n",
        "    return sentences, rel_label_ids, rel_arg_bindings, rel_labels"
      ],
      "metadata": {
        "id": "t9zLqP3wDb2v"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9hO0lZQ6Mvw_"
      },
      "source": [
        "## Data annotation\n",
        "Set of functions to properly annnotate the sentences using Bert reserved tokens"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "code_folding": [],
        "id": "Hj2tGdyBNuvg"
      },
      "outputs": [],
      "source": [
        "#Add entity markers to Bert vocabulary\n",
        "def update_vocab():\n",
        "    vocab = []\n",
        "\n",
        "    with open(BertREConfig.BERT_VOCAB_PATH, 'r') as F:\n",
        "        vocab = F.readlines()\n",
        "        vocab[BertREConfig.ENTITY1_START_TAG_ID] = BertREConfig.ENTITY1_START_TAG + \"\\n\"\n",
        "        vocab[BertREConfig.ENTITY1_END_TAG_ID] = BertREConfig.ENTITY1_END_TAG + \"\\n\"\n",
        "        vocab[BertREConfig.ENTITY2_START_TAG_ID] = BertREConfig.ENTITY2_START_TAG + \"\\n\"\n",
        "        vocab[BertREConfig.ENTITY2_END_TAG_ID] = BertREConfig.ENTITY2_END_TAG + \"\\n\"\n",
        "\n",
        "    with open(BertREConfig.BERT_VOCAB_PATH, \"w\") as F:\n",
        "        F.writelines(vocab)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "code_folding": [
          1
        ],
        "id": "12PkssBQNuvg"
      },
      "outputs": [],
      "source": [
        "\n",
        "#Tokenize sentence using Bert tokenizer, adding entity markers\n",
        "def tokenize_sentence(sentence, tokenizer, seq_length=BertREConfig.MAX_SEQ_LENGTH, is_test=False):\n",
        "    \n",
        "    tokens = [\"[CLS]\"]\n",
        "\n",
        "    entity_starts = []\n",
        "    entity_ends = []\n",
        "    \n",
        "    for token in tokenizer.tokenize(sentence)[:seq_length - 2]:\n",
        "        if token in [BertREConfig.ENTITY1_START_TAG, BertREConfig.ENTITY2_START_TAG]:\n",
        "            entity_starts.append(len(tokens))\n",
        "            \n",
        "        elif token in [BertREConfig.ENTITY1_END_TAG, BertREConfig.ENTITY2_END_TAG]:\n",
        "            entity_ends.append(len(tokens))\n",
        "            \n",
        "        tokens.append(token)\n",
        "    \n",
        "    tokens.append(\"[SEP]\")    \n",
        "        \n",
        "    if (len(entity_starts) != 2) or (len(entity_ends) != 2):\n",
        "        return False\n",
        "\n",
        "    if not is_test:\n",
        "        if np.random.rand() < BertREConfig.REPLACE_ARG_PROB:\n",
        "            e1_length_diff = (entity_ends[0] - entity_starts[0]) - 2\n",
        "\n",
        "            tokens = tokens[:entity_starts[0] + 1] + [\"[MASK]\"] + tokens[entity_ends[0]:]\n",
        "\n",
        "            entity_ends[0] = entity_starts[0] + 2        \n",
        "\n",
        "            entity_starts[1] = entity_starts[1] - e1_length_diff\n",
        "            entity_ends[1] = entity_ends[1] - e1_length_diff\n",
        "\n",
        "        if np.random.rand() < BertREConfig.REPLACE_ARG_PROB:\n",
        "            tokens = tokens[:entity_starts[1] + 1] + [\"[MASK]\"] + tokens[entity_ends[1]:]\n",
        "            entity_ends[1] = entity_starts[1] + 2        \n",
        "    \n",
        "    assert(tokens[entity_starts[0]] == BertREConfig.ENTITY1_START_TAG)\n",
        "    assert(tokens[entity_starts[1]] == BertREConfig.ENTITY2_START_TAG)\n",
        "    assert(tokens[entity_ends[0]] == BertREConfig.ENTITY1_END_TAG)\n",
        "    assert(tokens[entity_ends[1]] == BertREConfig.ENTITY2_END_TAG)\n",
        "\n",
        "    input_ids = tokenizer.convert_tokens_to_ids(tokens)\n",
        "        \n",
        "    return (input_ids, entity_starts[0], entity_starts[1], tokens)\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "def annotate_sentence(sentence, e1_begin, e1_end, e2_begin, e2_end):\n",
        "    \n",
        "    a1_start = min(e1_begin - 1, e2_begin)    \n",
        "    a1_end = min(e1_end + 1, e2_end + 1)\n",
        "\n",
        "    a2_start = max(e1_begin - 1, e2_begin)\n",
        "    a2_end = max(e1_end + 1, e2_end + 1)\n",
        "    \n",
        "    new_sentence = \" \".join([\n",
        "        sentence[:a1_start], \n",
        "        BertREConfig.ENTITY1_START_TAG, \n",
        "        sentence[a1_start:a1_end],\n",
        "        BertREConfig.ENTITY1_END_TAG, \n",
        "        sentence[a1_end:a2_start],\n",
        "        BertREConfig.ENTITY2_START_TAG, \n",
        "        sentence[a2_start:a2_end],\n",
        "        BertREConfig.ENTITY2_END_TAG, \n",
        "        sentence[a2_end:]\n",
        "    ])\n",
        "\n",
        "    return new_sentence  "
      ],
      "metadata": {
        "id": "csEKF-NtEzKX"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1ZuklXrSOGjr"
      },
      "source": [
        "## Feature Engineering\n",
        "RE Feature Engineering consists of token ids (input_ids), entities POS and label ids"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#Representation of RE featurues\n",
        "class REFeatures(object):\n",
        "    \n",
        "    def __str__(self):\n",
        "        return \"{} ({})\".format(\n",
        "            \", \".join(\n",
        "                map(lambda x: str(x), self.input_ids)), \n",
        "            self.sentence)\n",
        "    \n",
        "    def __init__(self,\n",
        "                 input_ids,\n",
        "                 entity1_pos,\n",
        "                 entity2_pos,\n",
        "                 rel_label_id,\n",
        "                 rel_arg_binding,\n",
        "                 sentence=\"\"):\n",
        "\n",
        "        self.input_ids = input_ids\n",
        "        self.entity1_pos = entity1_pos\n",
        "        self.entity2_pos = entity2_pos\n",
        "        self.rel_label_id = rel_label_id\n",
        "        self.rel_arg_binding = rel_arg_binding\n",
        "        self.sentence = sentence        \n"
      ],
      "metadata": {
        "id": "yomM8_4vEO0X"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#Create RE features from a list of sentencesand tarkets\n",
        "def make_features(sentences, targets, tokenizer, is_test=False):    \n",
        "    features = []\n",
        "    for i in range(len(sentences)):\n",
        "        ts = tokenize_sentence(sentences[i], tokenizer, is_test=is_test)\n",
        "        if ts:\n",
        "            features.append(\n",
        "                REFeatures(\n",
        "                    input_ids=ts[0],\n",
        "                    entity1_pos=ts[1],\n",
        "                    entity2_pos=ts[2],\n",
        "                    rel_label_id=targets[i][0],\n",
        "                    rel_arg_binding=targets[i][1],\n",
        "                    sentence=\" \".join(ts[3])\n",
        "                ))\n",
        "\n",
        "    return features"
      ],
      "metadata": {
        "id": "eSoQ9805ERqm"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4jtoNOaRTlq-"
      },
      "source": [
        "## Batches creation\n",
        "For feeding the training process"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "code_folding": [
          3
        ],
        "id": "7IShThYvNuvh"
      },
      "outputs": [],
      "source": [
        "#Make a batch of trainin/testing examples. \n",
        "#If max_seq_len is None, then use the sequence max length in the batch\n",
        "\n",
        "def make_batch(features, max_seq_len = None):\n",
        "    batch_size = len(features)\n",
        "    use_rel_args = BertREConfig.REL_ARG_BINDING_COLUMN is not None\n",
        "    if max_seq_len is None:\n",
        "        max_seq_len = max([len(f.input_ids) for f in features])\n",
        "    \n",
        "    input_ids = np.ones([batch_size, max_seq_len], dtype=np.int32) * BertREConfig.PAD_ID\n",
        "    input_mask = np.zeros([batch_size, max_seq_len], dtype=np.int32)\n",
        "    segment_ids = np.zeros([batch_size, max_seq_len], dtype=np.int32)\n",
        "    entity1_pos = np.zeros([batch_size], dtype=np.int32)\n",
        "    entity2_pos = np.zeros([batch_size], dtype=np.int32)\n",
        "    rel_label_ids = np.zeros([batch_size], dtype=np.int32)\n",
        "    rel_arg_bindings = np.zeros([batch_size], dtype=np.int32)\n",
        "    \n",
        "    i = 0\n",
        "    \n",
        "    for f in features:\n",
        "        \n",
        "        input_ids[i, :len(f.input_ids)] = np.array(f.input_ids)\n",
        "        input_mask[i, :len(f.input_ids)] = 1\n",
        "        rel_label_ids[i] = f.rel_label_id\n",
        "        rel_arg_bindings[i] = f.rel_arg_binding\n",
        "        entity1_pos[i] = f.entity1_pos\n",
        "        entity2_pos[i] = f.entity2_pos\n",
        "        i += 1\n",
        "    \n",
        "    batch = {\n",
        "        \"input_ids:0\": input_ids,\n",
        "        \"input_mask:0\": input_mask,\n",
        "        \"segment_ids:0\": segment_ids,\n",
        "        \"rel_label_ids:0\": rel_label_ids,\n",
        "        \"entity1_pos:0\": entity1_pos,\n",
        "        \"entity2_pos:0\": entity2_pos,\n",
        "    }\n",
        "    \n",
        "    if use_rel_args:\n",
        "        batch[\"rel_arg_bindings:0\"] = rel_arg_bindings\n",
        "        \n",
        "    return batch"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d9sho0T5S8EH"
      },
      "source": [
        "## Optimizer creation\n",
        "To carry out gradient descent and weight update with specific warm up, learning rate, etc."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "code_folding": [
          1
        ],
        "id": "N8RLqdQMNuvh"
      },
      "outputs": [],
      "source": [
        "#Create Bert RE optimizer graph\n",
        "def create_optimizer(loss, learning_rate, num_train_steps, num_warmup_steps, use_tpu):\n",
        "    \n",
        "    global_step = tf.train.get_or_create_global_step()\n",
        "\n",
        "    \n",
        "    \n",
        "    # Implements linear decay of the learning rate.\n",
        "    learning_rate = tf.train.polynomial_decay(\n",
        "      learning_rate,\n",
        "      global_step,\n",
        "      num_train_steps,\n",
        "      end_learning_rate=0.0,\n",
        "      power=1.0,\n",
        "      cycle=False)\n",
        "\n",
        "    tf.identity(learning_rate, name=\"c_lr\")\n",
        "    \n",
        "    # Implements linear warmup. I.e., if global_step < num_warmup_steps, the\n",
        "    # learning rate will be `global_step/num_warmup_steps * init_lr`.\n",
        "    \n",
        "    global_steps_int = tf.cast(global_step, tf.int32)\n",
        "    warmup_steps_int = num_warmup_steps\n",
        "\n",
        "    global_steps_float = tf.cast(global_steps_int, tf.float32)\n",
        "    warmup_steps_float = tf.cast(warmup_steps_int, tf.float32)\n",
        "\n",
        "    warmup_percent_done = global_steps_float / warmup_steps_float\n",
        "    warmup_learning_rate = learning_rate * warmup_percent_done\n",
        "\n",
        "    is_warmup = tf.cast(global_steps_int < warmup_steps_int, tf.float32)\n",
        "    learning_rate = (\n",
        "        (1.0 - is_warmup) * learning_rate + is_warmup * warmup_learning_rate)\n",
        "\n",
        "    \n",
        "    tf.identity(learning_rate, name=\"c_lr2\")\n",
        "    \n",
        "    # It is recommended that you use this optimizer for fine tuning, since this\n",
        "    # is how the model was trained (note that the Adam m/v variables are NOT\n",
        "    # loaded from init_checkpoint.)\n",
        "    optimizer = optimization.AdamWeightDecayOptimizer(\n",
        "      learning_rate=learning_rate,\n",
        "      weight_decay_rate=0.01,\n",
        "      beta_1=0.9,\n",
        "      beta_2=0.999,\n",
        "      epsilon=1e-6,\n",
        "      exclude_from_weight_decay=[\"LayerNorm\", \"layer_norm\", \"bias\"])\n",
        "\n",
        "    tvars = tf.trainable_variables()\n",
        "    grads = tf.gradients(loss, tvars)\n",
        "\n",
        "    # This is how the model was pre-trained.\n",
        "    (grads, _) = tf.clip_by_global_norm(grads, clip_norm=1.0)\n",
        "\n",
        "    train_op = optimizer.apply_gradients(\n",
        "      zip(grads, tvars), global_step=global_step)\n",
        "\n",
        "    # Normally the global step update is done inside of `apply_gradients`.\n",
        "    # However, `AdamWeightDecayOptimizer` doesn't do this. But if you use\n",
        "    # a different optimizer, you should probably take this line out.\n",
        "    new_global_step = global_step + 1\n",
        "    train_op = tf.group(train_op, [global_step.assign(new_global_step)], name=\"optimizer\")\n",
        "    return train_op\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Mo-ebbmCT1fR"
      },
      "source": [
        "## BERT model creation\n",
        "Creationg of the model using BERT architecture on TensorFlow"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#Create Bert RE model graph, for training (is_trainable = True) and for inference (is_trainable = False)\n",
        "def create_model(\n",
        "    num_relations, \n",
        "    num_arg_bindings=3, \n",
        "    num_hidden_units=BertREConfig.NUM_HIDDEN_UNITS, \n",
        "    chkpoint_path=None, \n",
        "    is_trainable=True):\n",
        "    \n",
        "    with tf.compat.v1.Session() as session:    \n",
        "\n",
        "        num_train_steps = tf.compat.v1.placeholder_with_default(\n",
        "            input=tf.constant(1000, dtype=tf.float32), shape=(), name=\"num_train_steps\")\n",
        "        \n",
        "        num_warm_up_steps = tf.compat.v1.placeholder_with_default(\n",
        "            input=tf.cast(tf.round(0.1 * num_train_steps), tf.int32), shape=(), name=\"num_warm_up_steps\")\n",
        "\n",
        "        input_ids = tf.compat.v1.placeholder(\n",
        "            dtype=tf.compat.v1.int32, shape=(None, None), name=\"input_ids\")\n",
        "        \n",
        "        batch_size = tf.shape(input_ids)[0]\n",
        "        seq_len = tf.shape(input_ids)[1]\n",
        "        \n",
        "        input_mask = tf.compat.v1.placeholder(\n",
        "            dtype=tf.compat.v1.int32, shape=(None, None), name=\"input_mask\")\n",
        "        \n",
        "        segment_ids = tf.compat.v1.placeholder(\n",
        "            dtype=tf.compat.v1.int32, shape=(None, None), name=\"segment_ids\")\n",
        "        \n",
        "        rel_label_ids = tf.compat.v1.placeholder(\n",
        "            dtype=tf.compat.v1.int32, shape=(None), name=\"rel_label_ids\")\n",
        "        \n",
        "        rel_arg_bindings = tf.compat.v1.placeholder_with_default(\n",
        "            input=tf.zeros(shape=(batch_size),dtype=tf.compat.v1.int32), shape=(None), name=\"rel_arg_bindings\")\n",
        "        \n",
        "        entity1_pos = tf.compat.v1.placeholder(\n",
        "            dtype=tf.compat.v1.int32, shape=(None), name=\"entity1_pos\")\n",
        "        \n",
        "        entity2_pos = tf.compat.v1.placeholder(\n",
        "            dtype=tf.compat.v1.int32, shape=(None), name=\"entity2_pos\")\n",
        "        \n",
        "        #set dropout to 0 if mode is not trainable\n",
        "        default_dropout_rate = BertREConfig.DROPOUT_RATE if is_trainable else 0.0        \n",
        "        dropout_rate = tf.compat.v1.placeholder_with_default(\n",
        "            input=tf.constant(default_dropout_rate, dtype=tf.float32), shape=(), name=\"dropout_rate\")\n",
        "        \n",
        "        learning_rate = tf.compat.v1.placeholder_with_default(\n",
        "            input=tf.constant(2e-5, dtype=tf.float32), shape=(), name=\"learning_rate\")\n",
        "        \n",
        "        config = modeling.BertConfig.from_json_file(BertREConfig.BERT_MODEL_CONFIG_PATH)\n",
        "\n",
        "#         breakpoint()\n",
        "        \n",
        "        bert_model = modeling.BertModel(\n",
        "            config=config,\n",
        "            is_training=is_trainable,\n",
        "            input_ids=input_ids,\n",
        "            input_mask=input_mask,\n",
        "            token_type_ids=segment_ids)\n",
        "\n",
        "        if chkpoint_path:\n",
        "            tvars = tf.trainable_variables()\n",
        "            (assignment_map, initialized_variable_names) = modeling.get_assignment_map_from_checkpoint(\n",
        "                tvars, chkpoint_path)\n",
        "            tf.train.init_from_checkpoint(chkpoint_path, assignment_map)    \n",
        "\n",
        "        output_layer = bert_model.get_sequence_output()\n",
        "\n",
        "        if BertREConfig.USE_ENTITY_POSITIONS:\n",
        "            #get entity start marker embeddings\n",
        "\n",
        "            #E1 mask\n",
        "            entity1_mask =  tf.repeat(\n",
        "                tf.one_hot(entity1_pos, seq_len), \n",
        "                config.hidden_size, \n",
        "                axis=1)\n",
        "\n",
        "            #E2 mask\n",
        "            entity2_mask =  tf.repeat(\n",
        "                tf.one_hot(entity2_pos, seq_len), \n",
        "                config.hidden_size, \n",
        "                axis=1)\n",
        "\n",
        "\n",
        "            #Hidden layer representation for E1\n",
        "            entity1_embd = tf.reduce_sum(\n",
        "                tf.reshape(\n",
        "                    (tf.reshape(output_layer, shape=(batch_size, -1)) * entity1_mask), \n",
        "                    shape=[batch_size, seq_len, config.hidden_size]), \n",
        "                axis=1)\n",
        "\n",
        "            #Hidden layer representation for E2\n",
        "            entity2_embd = tf.reduce_sum(\n",
        "                tf.reshape(\n",
        "                    (tf.reshape(output_layer, shape=(batch_size, -1)) * entity2_mask), \n",
        "                    shape=[batch_size, seq_len, config.hidden_size]), \n",
        "                axis=1)\n",
        "\n",
        "            #Concat representions\n",
        "            if BertREConfig.USE_CLS_POSITION:                \n",
        "                classification_layer = tf.concat([entity1_embd, entity2_embd, output_layer[:,0,:]], axis=1)\n",
        "            else:\n",
        "                classification_layer = tf.concat([entity1_embd, entity2_embd], axis=1)\n",
        "        else:\n",
        "            if not BertREConfig.USE_CLS_POSITION:\n",
        "                raise(\"Either USE_ENTITY_POSITIONS or USE_CLS_POSITION should be set to True.\")\n",
        "            else:\n",
        "                classification_layer = output_layer[:,0,:]\n",
        "                \n",
        "#         print(\"Classification layer size: \")\n",
        "#         print(classification_layer.shape)\n",
        "        '''Add full connection layer and dropout layer'''\n",
        "        \n",
        "        if num_hidden_units > 0:\n",
        "            fc = tf.layers.dense(\n",
        "                classification_layer, \n",
        "                num_hidden_units, \n",
        "                name='fc1')\n",
        "            fc = tf.nn.relu(fc)\n",
        "        else:\n",
        "            fc = tf.identity(classification_layer, name='fc1')\n",
        "\n",
        "        fc = tf.nn.dropout(fc, rate=dropout_rate)\n",
        "\n",
        "        '''logits'''\n",
        "        rel_label_logits = tf.layers.dense(fc, num_relations, name='rel_label_logits')\n",
        "        rel_label_log_probs = tf.nn.softmax(rel_label_logits, name=\"rel_label_probs\")\n",
        "        rel_label_predictions = tf.argmax(\n",
        "            rel_label_log_probs, axis=-1, output_type=tf.int32, name=\"rel_label_predictions\")\n",
        "\n",
        "        if num_arg_bindings > 1:    \n",
        "            rel_arg_binding_logits = tf.layers.dense(fc, num_arg_bindings, name='rel_arg_binding_logits')\n",
        "            rel_arg_binding_probs = tf.nn.softmax(rel_arg_binding_logits, name=\"rel_arg_binding_probs\")\n",
        "            rel_arg_binding_predictions = tf.argmax(\n",
        "                rel_arg_binding_probs, axis=-1, output_type=tf.int32, name=\"rel_arg_binding_predictions\")\n",
        "        else:\n",
        "            rel_arg_binding_probs = tf.ones_like(rel_arg_bindings, name=\"rel_arg_binding_probs\")\n",
        "            rel_arg_binding_predictions = tf.zeros_like(\n",
        "                rel_arg_bindings, name=\"rel_arg_binding_predictions\")\n",
        "\n",
        "        '''Calculate loss. Convert predicted labels into one hot form. '''            \n",
        "        rel_label_targets = tf.one_hot(rel_label_ids, depth=num_relations)\n",
        "        rel_label_loss = tf.nn.softmax_cross_entropy_with_logits_v2(\n",
        "                labels=rel_label_targets,\n",
        "                logits=rel_label_logits)\n",
        "\n",
        "        if num_arg_bindings > 1:                \n",
        "            rel_arg_binding_targets = tf.one_hot(rel_arg_bindings, depth=num_arg_bindings)\n",
        "            rel_arg_binding_loss = tf.nn.softmax_cross_entropy_with_logits_v2(\n",
        "                    labels=rel_arg_binding_targets,\n",
        "                    logits=rel_arg_binding_logits)\n",
        "        else:\n",
        "            rel_arg_binding_loss = 0                    \n",
        "        \n",
        "        rel_label_example_accuracy = tf.cast(\n",
        "            tf.equal(rel_label_predictions, rel_label_ids), tf.float32, name=\"rel_label_acc\")\n",
        "        \n",
        "        rel_label_accuracy = tf.reduce_mean(rel_label_example_accuracy, name=\"rel_label_mean_acc\")        \n",
        "        \n",
        "        rel_arg_binding_example_accuracy = tf.cast(\n",
        "            tf.equal(rel_arg_binding_predictions, rel_arg_bindings), tf.float32, name=\"rel_arg_binding_acc\")\n",
        "        \n",
        "        rel_arg_binding_accuracy = tf.reduce_mean(\n",
        "            rel_arg_binding_example_accuracy, name=\"rel_arg_binding_mean_acc\")\n",
        "        \n",
        "        total_example_accuracy = tf.identity(\n",
        "            rel_label_example_accuracy * rel_arg_binding_example_accuracy, name=\"total_acc\")\n",
        "        \n",
        "        total_accuracy = tf.identity(total_example_accuracy, name=\"total_mean_acc\")\n",
        "        \n",
        "        \n",
        "        loss = tf.reduce_mean(rel_label_loss + rel_arg_binding_loss, name=\"loss\")\n",
        "\n",
        "        if is_trainable:\n",
        "            train_op = create_optimizer(\n",
        "                loss, \n",
        "                learning_rate, \n",
        "                num_train_steps, \n",
        "                num_warm_up_steps,\n",
        "                use_tpu=False)        \n",
        "        else:\n",
        "            train_op = tf.no_op()\n",
        "            \n",
        "        init = tf.global_variables_initializer()\n",
        "\n",
        "        return (\n",
        "                train_op,\n",
        "                loss,       \n",
        "                total_accuracy,\n",
        "                (rel_label_predictions, rel_arg_binding_predictions), \n",
        "                (rel_label_log_probs, rel_arg_binding_probs))"
      ],
      "metadata": {
        "id": "RehcGAB_FPTW"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2XWjsPywUsNd"
      },
      "source": [
        "## Model saving\n",
        "Function to export the trained BERT model to disk in TF"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "code_folding": [],
        "id": "g8SMo6jvNuvh"
      },
      "outputs": [],
      "source": [
        "def export_model(model_id, is_trainable = True, num_arg_bindings = 3):\n",
        "    \n",
        "    with tf.compat.v1.Session() as session:\n",
        "    \n",
        "        model = create_model(\n",
        "            len(rel_labels), \n",
        "            is_trainable=is_trainable,\n",
        "            num_arg_bindings=num_arg_bindings)\n",
        "    \n",
        "        input_tensors = {}\n",
        "        output_tensors = {}\n",
        "\n",
        "        input_tensors_names = {\n",
        "            \"input_ids:0\",\n",
        "            \"input_mask:0\",\n",
        "            \"segment_ids:0\",\n",
        "            \"entity1_pos:0\",\n",
        "            \"entity2_pos:0\",\n",
        "        }\n",
        "\n",
        "        output_tensors_names = [\n",
        "            \"loss:0\",\n",
        "            \"rel_label_acc:0\",\n",
        "            \"rel_arg_binding_acc:0\",\n",
        "            \"total_acc:0\",\n",
        "            \"rel_label_probs:0\",\n",
        "            \"rel_label_predictions:0\",\n",
        "            \"rel_arg_binding_probs:0\",\n",
        "            \"rel_arg_binding_predictions:0\"        \n",
        "        ]\n",
        "\n",
        "\n",
        "        for k in input_tensors_names:\n",
        "            t = session.graph.get_tensor_by_name(k)\n",
        "            input_tensors[t.name] = t\n",
        "\n",
        "        for k in output_tensors_names:\n",
        "            t = session.graph.get_tensor_by_name(k)\n",
        "            output_tensors[k] = t\n",
        "        \n",
        "        # print(\"{} trainable variables: \".format(len(tf.trainable_variables())))\n",
        "        size_f = lambda v: reduce(lambda x, y: x*y, v.get_shape().as_list())\n",
        "        n = sum(size_f(v) for v in tf.trainable_variables())\n",
        "        print(\"{} trainbale parameters.\".format(n))    \n",
        "        \n",
        "        tf.train.Saver().restore(session, f\"{BertREConfig.CHKPOINT_PATH}/{model_id}/model\")\n",
        "\n",
        "        shutil.rmtree(BertREConfig.EXPORT_PATH, ignore_errors=True)\n",
        "\n",
        "        #save model\n",
        "        tf.saved_model.simple_save(\n",
        "            session,\n",
        "            BertREConfig.EXPORT_PATH,\n",
        "            inputs=input_tensors,\n",
        "            outputs=output_tensors\n",
        "        )\n",
        "\n",
        "        #copy assets to the destiation folder\n",
        "        shutil.copytree(\n",
        "            f\"{BertREConfig.CHKPOINT_PATH}/{model_id}/assets\", \n",
        "            f\"{BertREConfig.EXPORT_PATH}/assets\")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CcIgaAyfYH_V"
      },
      "source": [
        "## Training\n",
        "Function to train the model for RE"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "code_folding": [],
        "id": "SPzjif4FNuvi"
      },
      "outputs": [],
      "source": [
        "#Train a Bert RE model and save it in the checkpoints folder\n",
        "def train_model(model_id, train_features, test_features, rel_labels, num_arg_bindings = 3):\n",
        "\n",
        "    ops.reset_default_graph()\n",
        "    \n",
        "    use_rel_args = BertREConfig.REL_ARG_BINDING_COLUMN is not None\n",
        "    \n",
        "    with tf.compat.v1.Session() as session:\n",
        "\n",
        "        model = create_model(\n",
        "            len(rel_labels), \n",
        "            chkpoint_path=BertREConfig.BERT_MODEL_CHECKPOINT_PATH, \n",
        "            num_arg_bindings=num_arg_bindings, \n",
        "            is_trainable=True\n",
        "        )\n",
        "        # return\n",
        "        session.run(\"init\")\n",
        "        \n",
        "        num_train_steps = (BertREConfig.NUM_TRAIN_EPOCHS * len(train_features)) // BertREConfig.BATCH_SIZE    \n",
        "\n",
        "        clear_output(wait=True)\n",
        "\n",
        "        print(\"{:^11}{:^11}{:>10}{:>10}{:>10}{:>10}{:>10}{:>10}{:>10}\".format(\n",
        "            \"Epoch\", \"Batch\", \n",
        "            \"Loss\", \n",
        "            \"L_ACC\", \"Arg_ACC\", \"ACC\",\n",
        "            \"vL_ACC\", \"vArg_ACC\", \"vACC\"))\n",
        "\n",
        "        for e in range(BertREConfig.NUM_TRAIN_EPOCHS):\n",
        "\n",
        "            np.random.shuffle(train_features)            \n",
        "\n",
        "            b_loss = []\n",
        "            b_rel_label_acc = []\n",
        "            b_rel_arg_binding_acc = []\n",
        "            b_total_acc = []\n",
        "            for b in range(0, len(train_features) // BertREConfig.BATCH_SIZE):\n",
        "\n",
        "                batch = make_batch(\n",
        "                    train_features[b * BertREConfig.BATCH_SIZE: (b + 1) * BertREConfig.BATCH_SIZE]\n",
        "                )#, max_seq_len=MAX_SEQ_LENGTH)\n",
        "\n",
        "                data = batch\n",
        "\n",
        "                if b == 0:\n",
        "                    data[\"num_train_steps:0\"] = num_train_steps\n",
        "                    data[\"learning_rate:0\"] = BertREConfig.LEARNING_RATE\n",
        "\n",
        "                eval_tensors = [\n",
        "                    \"optimizer\", \n",
        "                    \"loss:0\", \n",
        "                    \"rel_label_mean_acc:0\",\n",
        "                    \"rel_arg_binding_mean_acc:0\",\n",
        "                    \"total_mean_acc:0\"\n",
        "                ]\n",
        "                _, loss, rel_label_acc, rel_arg_bind_acc, total_acc = session.run(\n",
        "                    eval_tensors, feed_dict=data)\n",
        "                b_loss.append(loss)\n",
        "                b_rel_label_acc.append(rel_label_acc)\n",
        "                b_rel_arg_binding_acc.append(rel_arg_bind_acc)\n",
        "                b_total_acc.append(total_acc)\n",
        "\n",
        "                print(\"\\r{:>5}/{:<5}{:>5}/{:<5}{:>10.4f}{:>10.3f}{:>10.3f}{:>10.3f}\".format(\n",
        "                        e+1,\n",
        "                        BertREConfig.NUM_TRAIN_EPOCHS,\n",
        "                        b + 1, \n",
        "                        len(train_features) // BertREConfig.BATCH_SIZE,\n",
        "                        np.mean(b_loss), \n",
        "                        np.mean(b_rel_label_acc), \n",
        "                        np.mean(b_rel_arg_binding_acc), \n",
        "                        np.mean(b_total_acc)  \n",
        "                    ), end=\"\")\n",
        "\n",
        "\n",
        "            v_rel_label_acc = []\n",
        "            v_rel_arg_binding_acc = []\n",
        "            v_total_acc = []\n",
        "\n",
        "            for v_b in range(0, len(test_features) // BertREConfig.V_BATCH_SIZE):\n",
        "                batch = make_batch(\n",
        "                    test_features[v_b * BertREConfig.V_BATCH_SIZE: (v_b + 1) * BertREConfig.V_BATCH_SIZE])\n",
        "\n",
        "                data = batch\n",
        "\n",
        "                eval_tensors = [\n",
        "                    \"rel_label_mean_acc:0\",\n",
        "                    \"rel_arg_binding_mean_acc:0\",\n",
        "                    \"total_mean_acc:0\",\n",
        "                ]\n",
        "\n",
        "                rel_label_acc, rel_arg_bind_acc, total_acc = session.run(eval_tensors, feed_dict=data)\n",
        "                v_rel_label_acc.append(rel_label_acc)\n",
        "                v_rel_arg_binding_acc.append(rel_arg_bind_acc)\n",
        "                v_total_acc.append(total_acc)\n",
        "\n",
        "            print(\"{:>10.3f}{:>10.3f}{:>10.3f}\".format(\n",
        "                np.mean(v_rel_label_acc), \n",
        "                np.mean(v_rel_arg_binding_acc), \n",
        "                np.mean(v_total_acc)))                \n",
        "\n",
        "\n",
        "        \n",
        "        shutil.rmtree(f\"{BertREConfig.CHKPOINT_PATH}/{model_id}\", ignore_errors=True)\n",
        "        os.mkdir(f\"{BertREConfig.CHKPOINT_PATH}/{model_id}\")\n",
        "        \n",
        "        saver = tf.train.Saver()\n",
        "        saver.save(session, f\"{BertREConfig.CHKPOINT_PATH}/{model_id}/model\")\n",
        "        \n",
        "        os.mkdir(f\"{BertREConfig.CHKPOINT_PATH}/{model_id}/assets/\")\n",
        "        \n",
        "        shutil.copy(\n",
        "            BertREConfig.BERT_VOCAB_PATH, \n",
        "            f\"{BertREConfig.CHKPOINT_PATH}/{model_id}/assets/vocab.txt\")\n",
        "        \n",
        "        with open(f\"{BertREConfig.CHKPOINT_PATH}/{model_id}/assets/categories.txt\", \"wt\") as F:\n",
        "            F.writelines(\"\\n\".join(rel_labels))\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Sfy0E-sdYJmV"
      },
      "source": [
        "## Evaluation\n",
        "Functions to get the metrics on the RE model and print them"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "code_folding": [
          62
        ],
        "id": "zoxJlAU1Nuvi"
      },
      "outputs": [],
      "source": [
        "def eval_metrics(model_id, features, rel_labels, num_arg_bindings = 3, exclude_rels=[]):\n",
        "    with tf.compat.v1.Session() as session:\n",
        "\n",
        "        model = create_model(\n",
        "            len(rel_labels), \n",
        "            num_arg_bindings=num_arg_bindings,\n",
        "            is_trainable=False)\n",
        "        \n",
        "        tf.train.Saver().restore(session, f\"{BertREConfig.CHKPOINT_PATH}/{model_id}/model\")\n",
        "\n",
        "        metrics_data = {}\n",
        "        for rel in rel_labels:\n",
        "            metrics_data[rel] = ([], [], [])\n",
        "            \n",
        "        for v_b in range(0, len(features) // BertREConfig.V_BATCH_SIZE):\n",
        "            batch = make_batch(\n",
        "                features[v_b * BertREConfig.V_BATCH_SIZE: (v_b + 1) * BertREConfig.V_BATCH_SIZE])\n",
        "\n",
        "            data = batch\n",
        "\n",
        "            eval_tensors = [\"total_acc:0\", \"rel_label_ids:0\", \"rel_label_predictions:0\"]\n",
        "\n",
        "            total_acc, rel_label_ids, rel_label_preds = session.run(eval_tensors, feed_dict=data)\n",
        "            \n",
        "            for i in range(len(rel_label_ids)):\n",
        "                acc = total_acc[i]\n",
        "                pred = rel_label_preds[i]\n",
        "                target = rel_label_ids[i]\n",
        "                rel_target = rel_labels[target]\n",
        "                rel_pred = rel_labels[pred]\n",
        "                \n",
        "                metrics_data[rel_target][2].append(1)\n",
        "                \n",
        "                if acc:\n",
        "                    metrics_data[rel_target][0].append(1)\n",
        "                    metrics_data[rel_pred][1].append(1)\n",
        "                else:\n",
        "                    metrics_data[rel_target][0].append(0)\n",
        "                    metrics_data[rel_pred][1].append(0)                \n",
        "\n",
        "        results = {}        \n",
        "        \n",
        "        for rel in [rel for rel in rel_labels if rel not in exclude_rels]:\n",
        "            if len(metrics_data[rel][0]):\n",
        "                recall = np.mean(metrics_data[rel][0])\n",
        "            else:\n",
        "                recall = 0\n",
        "            if len(metrics_data[rel][1]):\n",
        "                precision = np.mean(metrics_data[rel][1])\n",
        "            else:\n",
        "                precision = 0\n",
        "            if (recall + precision):\n",
        "                f1 = 2 * (recall * precision) / (recall + precision)\n",
        "            else:\n",
        "                f1 = np.NaN\n",
        "               \n",
        "            support = np.sum(metrics_data[rel][2])\n",
        "            \n",
        "            results[rel] = (recall, precision, f1, support)\n",
        "        \n",
        "        return results\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "def print_metrics(results):\n",
        "    print(\"\\n\")\n",
        "    print(\"{:<15}{:>10}{:>10}{:>10}{:>10}\\n\".format(\"Relation\", \"Recall\", \"Precision\", \"F1\", \"Support\"))\n",
        "\n",
        "    for rel in results:\n",
        "\n",
        "        print(f\"{rel:<15}{results[rel][0]:>10.3f}{results[rel][1]:>10.3f}{results[rel][2]:>10.3f}{results[rel][3]:>10}\")\n",
        "\n",
        "    mean_recall = np.mean([results[rel][0] for rel in results])\n",
        "    mean_precision = np.mean([results[rel][1] for rel in results])\n",
        "    mean_f1 = np.mean([results[rel][2] for rel in results])\n",
        "\n",
        "    support_sum = np.sum([results[rel][3] for rel in results])\n",
        "\n",
        "    w_mean_recall = np.sum([results[rel][0] * results[rel][3] for rel in results]) / support_sum\n",
        "    w_mean_precision = np.sum([results[rel][1] * results[rel][3] for rel in results]) / support_sum\n",
        "    w_mean_f1 = np.sum([results[rel][2] * results[rel][3] for rel in results]) / support_sum\n",
        "\n",
        "\n",
        "    metrics_name = \"Avg.\"\n",
        "\n",
        "    print(f\"\\n{metrics_name:<15}{mean_recall:>10.3f}{mean_precision:>10.3f}{mean_f1:>10.3f}\")\n",
        "\n",
        "    metrics_name = \"Weighted Avg.\"\n",
        "\n",
        "    print(f\"\\n{metrics_name:<15}{w_mean_recall:>10.3f}{w_mean_precision:>10.3f}{w_mean_f1:>10.3f}\")"
      ],
      "metadata": {
        "id": "Vz6KOv__GMfv"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "8vMnRd-fGbKf"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5nW1WF2UYhZ6"
      },
      "source": [
        "## MAIN: STEP-BY-STEP RE MODEL TRAINING EXECUTION"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6l_Z2MvCZHTs"
      },
      "source": [
        "### Update Bert vocabulary with special tokens"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#Update Bert vocabylary\n",
        "update_vocab()"
      ],
      "metadata": {
        "id": "VuP4FyPqG3rW"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dM26YYEtaJao"
      },
      "source": [
        "### Creating a Tokenizer"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#create tokenizer\n",
        "tokenizer = tokenization.FullTokenizer(vocab_file=BertREConfig.BERT_VOCAB_PATH, do_lower_case=False)"
      ],
      "metadata": {
        "id": "0A63xkdNG3wV"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wWtxzECjbSIb"
      },
      "source": [
        "### Read the input data.\n",
        "It should look like as follows (see output) and have the following columns"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "guu863lUfmT7"
      },
      "source": [
        "| Column      |          Explanation                     |\n",
        "|:-----------:|:----------------------------------------:|\n",
        "|dataset      | train/test                               |\n",
        "|source       | data provider                            |\n",
        "|txt_file     | .txt file                                |\n",
        "|sentence     | tokenized text sentence                  |\n",
        "|sent_id      | sentence id                              |\n",
        "|chunk1       | first entity                             |\n",
        "|begin1       | first token number of the first entity   |\n",
        "|end1         | last token number of the first entity    |\n",
        "|rel          | relation (O for no-relation)             |\n",
        "|chunk2       | second entity                            |\n",
        "|begin2       | first token number of the second entity  |\n",
        "|end2         | last token number of the second entity   |\n",
        "|label1       | label of the first entity                |\n",
        "|label2       | label of the second entity               |\n",
        "|lastCharEnt1 | last char number of the first entity     |\n",
        "|firstCharEnt1| first char number of the first entity    |\n",
        "|lastCharEnt2 | last char number of the second entity    |\n",
        "|firstCharEnt2| first char number of the second entity   |\n",
        "|words_in_ent1| number of words in first entity          |\n",
        "|words_in_ent2| number of words in second entity         |\n",
        "|words_between| word between entities                    |\n",
        "|is_train     | is it used for training?                 |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GAMHT9Cki-G9"
      },
      "source": [
        "Yes, We are ready to train REDL model. Now we will train a REDL model to get relations between **DOC**, **PARTY**, **ALIAS** and **EFFDATE** entitties. \n",
        "\n",
        "Let's look at our dataset. \n",
        "\n",
        "Your dataset have to be like this format."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "tags": [],
        "id": "SHLivW7TNuvi",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 424
        },
        "outputId": "625bfaaa-7834-4b87-dc33-8506d91ff61d"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "                                                   text  firstCharEnt1  \\\n",
              "0     EXHIBIT 10.43 Dated 29/3/18\\n\\nDistributorship...             29   \n",
              "1     EXHIBIT 10.43 Dated 29/3/18\\n\\nDistributorship...             29   \n",
              "2     EXHIBIT 10.43 Dated 29/3/18\\n\\nDistributorship...             20   \n",
              "3     Sections 200.80(b)(4) and Rule 406 of the Secu...            173   \n",
              "4     Monrovia, CA 91016 USA\\n\\n(hereinafter called ...             63   \n",
              "...                                                 ...            ...   \n",
              "3496  By execution of this Supplier/Subcontractor Co...             85   \n",
              "3497    /s Liu Gang Name: LIU GANG   Title: Authoriz...             20   \n",
              "3498  HOFV: HOF VILLAGE, LLC By: /s / Brian Parisi N...            193   \n",
              "3499  By: /s/ Robert Mattacchione Name: Robert Matta...             34   \n",
              "3500  EXHIBIT 10.2 EXECUTION VERSION NON-COMPETITION...            591   \n",
              "\n",
              "      firstCharEnt2  lastCharEnt1  lastCharEnt2  \\\n",
              "0                65            54            95   \n",
              "1               102            54           129   \n",
              "2                29            27            54   \n",
              "3               236           196           247   \n",
              "4               170            97           173   \n",
              "...             ...           ...           ...   \n",
              "3496             85            93            93   \n",
              "3497             38            28            58   \n",
              "3498            212           204           227   \n",
              "3499             61            53            64   \n",
              "3500            609           604           627   \n",
              "\n",
              "                                  chunk1                          chunk2  \\\n",
              "0              Distributorship agreement  Signature Orthopaedics Pty Ltd   \n",
              "1              Distributorship agreement     CPM Medical Consultants LLC   \n",
              "2                                29/3/18       Distributorship agreement   \n",
              "3                Collaboration Agreement                     Xencor, Inc   \n",
              "4     Boehringer Ingelheim International                             BII   \n",
              "...                                  ...                             ...   \n",
              "3496                            Supplier                        Supplier   \n",
              "3497                            LIU GANG            Authorized Signatory   \n",
              "3498                         David Baker                 President & CEO   \n",
              "3499                 Robert Mattacchione                             CEO   \n",
              "3500                       Gulf Houghton              Gulf International   \n",
              "\n",
              "              label1         label2        rel  direction  \n",
              "0                DOC          PARTY  signed_by          1  \n",
              "1                DOC          PARTY  signed_by          1  \n",
              "2            EFFDATE            DOC   dated_as          2  \n",
              "3                DOC          PARTY  signed_by          1  \n",
              "4              PARTY          ALIAS  has_alias          1  \n",
              "...              ...            ...        ...        ...  \n",
              "3496           PARTY           ROLE      other          0  \n",
              "3497  SIGNING_PERSON  SIGNING_TITLE      other          0  \n",
              "3498  SIGNING_PERSON  SIGNING_TITLE      other          0  \n",
              "3499  SIGNING_PERSON  SIGNING_TITLE      other          0  \n",
              "3500             ORG            ORG      other          0  \n",
              "\n",
              "[3501 rows x 11 columns]"
            ],
            "text/html": [
              "\n",
              "  <div id=\"df-667a7ab1-5964-45d2-ab73-627905bfd56f\">\n",
              "    <div class=\"colab-df-container\">\n",
              "      <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>text</th>\n",
              "      <th>firstCharEnt1</th>\n",
              "      <th>firstCharEnt2</th>\n",
              "      <th>lastCharEnt1</th>\n",
              "      <th>lastCharEnt2</th>\n",
              "      <th>chunk1</th>\n",
              "      <th>chunk2</th>\n",
              "      <th>label1</th>\n",
              "      <th>label2</th>\n",
              "      <th>rel</th>\n",
              "      <th>direction</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>0</th>\n",
              "      <td>EXHIBIT 10.43 Dated 29/3/18\\n\\nDistributorship...</td>\n",
              "      <td>29</td>\n",
              "      <td>65</td>\n",
              "      <td>54</td>\n",
              "      <td>95</td>\n",
              "      <td>Distributorship agreement</td>\n",
              "      <td>Signature Orthopaedics Pty Ltd</td>\n",
              "      <td>DOC</td>\n",
              "      <td>PARTY</td>\n",
              "      <td>signed_by</td>\n",
              "      <td>1</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>1</th>\n",
              "      <td>EXHIBIT 10.43 Dated 29/3/18\\n\\nDistributorship...</td>\n",
              "      <td>29</td>\n",
              "      <td>102</td>\n",
              "      <td>54</td>\n",
              "      <td>129</td>\n",
              "      <td>Distributorship agreement</td>\n",
              "      <td>CPM Medical Consultants LLC</td>\n",
              "      <td>DOC</td>\n",
              "      <td>PARTY</td>\n",
              "      <td>signed_by</td>\n",
              "      <td>1</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>2</th>\n",
              "      <td>EXHIBIT 10.43 Dated 29/3/18\\n\\nDistributorship...</td>\n",
              "      <td>20</td>\n",
              "      <td>29</td>\n",
              "      <td>27</td>\n",
              "      <td>54</td>\n",
              "      <td>29/3/18</td>\n",
              "      <td>Distributorship agreement</td>\n",
              "      <td>EFFDATE</td>\n",
              "      <td>DOC</td>\n",
              "      <td>dated_as</td>\n",
              "      <td>2</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>3</th>\n",
              "      <td>Sections 200.80(b)(4) and Rule 406 of the Secu...</td>\n",
              "      <td>173</td>\n",
              "      <td>236</td>\n",
              "      <td>196</td>\n",
              "      <td>247</td>\n",
              "      <td>Collaboration Agreement</td>\n",
              "      <td>Xencor, Inc</td>\n",
              "      <td>DOC</td>\n",
              "      <td>PARTY</td>\n",
              "      <td>signed_by</td>\n",
              "      <td>1</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>4</th>\n",
              "      <td>Monrovia, CA 91016 USA\\n\\n(hereinafter called ...</td>\n",
              "      <td>63</td>\n",
              "      <td>170</td>\n",
              "      <td>97</td>\n",
              "      <td>173</td>\n",
              "      <td>Boehringer Ingelheim International</td>\n",
              "      <td>BII</td>\n",
              "      <td>PARTY</td>\n",
              "      <td>ALIAS</td>\n",
              "      <td>has_alias</td>\n",
              "      <td>1</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>...</th>\n",
              "      <td>...</td>\n",
              "      <td>...</td>\n",
              "      <td>...</td>\n",
              "      <td>...</td>\n",
              "      <td>...</td>\n",
              "      <td>...</td>\n",
              "      <td>...</td>\n",
              "      <td>...</td>\n",
              "      <td>...</td>\n",
              "      <td>...</td>\n",
              "      <td>...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>3496</th>\n",
              "      <td>By execution of this Supplier/Subcontractor Co...</td>\n",
              "      <td>85</td>\n",
              "      <td>85</td>\n",
              "      <td>93</td>\n",
              "      <td>93</td>\n",
              "      <td>Supplier</td>\n",
              "      <td>Supplier</td>\n",
              "      <td>PARTY</td>\n",
              "      <td>ROLE</td>\n",
              "      <td>other</td>\n",
              "      <td>0</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>3497</th>\n",
              "      <td>/s Liu Gang Name: LIU GANG   Title: Authoriz...</td>\n",
              "      <td>20</td>\n",
              "      <td>38</td>\n",
              "      <td>28</td>\n",
              "      <td>58</td>\n",
              "      <td>LIU GANG</td>\n",
              "      <td>Authorized Signatory</td>\n",
              "      <td>SIGNING_PERSON</td>\n",
              "      <td>SIGNING_TITLE</td>\n",
              "      <td>other</td>\n",
              "      <td>0</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>3498</th>\n",
              "      <td>HOFV: HOF VILLAGE, LLC By: /s / Brian Parisi N...</td>\n",
              "      <td>193</td>\n",
              "      <td>212</td>\n",
              "      <td>204</td>\n",
              "      <td>227</td>\n",
              "      <td>David Baker</td>\n",
              "      <td>President &amp; CEO</td>\n",
              "      <td>SIGNING_PERSON</td>\n",
              "      <td>SIGNING_TITLE</td>\n",
              "      <td>other</td>\n",
              "      <td>0</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>3499</th>\n",
              "      <td>By: /s/ Robert Mattacchione Name: Robert Matta...</td>\n",
              "      <td>34</td>\n",
              "      <td>61</td>\n",
              "      <td>53</td>\n",
              "      <td>64</td>\n",
              "      <td>Robert Mattacchione</td>\n",
              "      <td>CEO</td>\n",
              "      <td>SIGNING_PERSON</td>\n",
              "      <td>SIGNING_TITLE</td>\n",
              "      <td>other</td>\n",
              "      <td>0</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>3500</th>\n",
              "      <td>EXHIBIT 10.2 EXECUTION VERSION NON-COMPETITION...</td>\n",
              "      <td>591</td>\n",
              "      <td>609</td>\n",
              "      <td>604</td>\n",
              "      <td>627</td>\n",
              "      <td>Gulf Houghton</td>\n",
              "      <td>Gulf International</td>\n",
              "      <td>ORG</td>\n",
              "      <td>ORG</td>\n",
              "      <td>other</td>\n",
              "      <td>0</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "<p>3501 rows × 11 columns</p>\n",
              "</div>\n",
              "      <button class=\"colab-df-convert\" onclick=\"convertToInteractive('df-667a7ab1-5964-45d2-ab73-627905bfd56f')\"\n",
              "              title=\"Convert this dataframe to an interactive table.\"\n",
              "              style=\"display:none;\">\n",
              "        \n",
              "  <svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\"viewBox=\"0 0 24 24\"\n",
              "       width=\"24px\">\n",
              "    <path d=\"M0 0h24v24H0V0z\" fill=\"none\"/>\n",
              "    <path d=\"M18.56 5.44l.94 2.06.94-2.06 2.06-.94-2.06-.94-.94-2.06-.94 2.06-2.06.94zm-11 1L8.5 8.5l.94-2.06 2.06-.94-2.06-.94L8.5 2.5l-.94 2.06-2.06.94zm10 10l.94 2.06.94-2.06 2.06-.94-2.06-.94-.94-2.06-.94 2.06-2.06.94z\"/><path d=\"M17.41 7.96l-1.37-1.37c-.4-.4-.92-.59-1.43-.59-.52 0-1.04.2-1.43.59L10.3 9.45l-7.72 7.72c-.78.78-.78 2.05 0 2.83L4 21.41c.39.39.9.59 1.41.59.51 0 1.02-.2 1.41-.59l7.78-7.78 2.81-2.81c.8-.78.8-2.07 0-2.86zM5.41 20L4 18.59l7.72-7.72 1.47 1.35L5.41 20z\"/>\n",
              "  </svg>\n",
              "      </button>\n",
              "      \n",
              "  <style>\n",
              "    .colab-df-container {\n",
              "      display:flex;\n",
              "      flex-wrap:wrap;\n",
              "      gap: 12px;\n",
              "    }\n",
              "\n",
              "    .colab-df-convert {\n",
              "      background-color: #E8F0FE;\n",
              "      border: none;\n",
              "      border-radius: 50%;\n",
              "      cursor: pointer;\n",
              "      display: none;\n",
              "      fill: #1967D2;\n",
              "      height: 32px;\n",
              "      padding: 0 0 0 0;\n",
              "      width: 32px;\n",
              "    }\n",
              "\n",
              "    .colab-df-convert:hover {\n",
              "      background-color: #E2EBFA;\n",
              "      box-shadow: 0px 1px 2px rgba(60, 64, 67, 0.3), 0px 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
              "      fill: #174EA6;\n",
              "    }\n",
              "\n",
              "    [theme=dark] .colab-df-convert {\n",
              "      background-color: #3B4455;\n",
              "      fill: #D2E3FC;\n",
              "    }\n",
              "\n",
              "    [theme=dark] .colab-df-convert:hover {\n",
              "      background-color: #434B5C;\n",
              "      box-shadow: 0px 1px 3px 1px rgba(0, 0, 0, 0.15);\n",
              "      filter: drop-shadow(0px 1px 2px rgba(0, 0, 0, 0.3));\n",
              "      fill: #FFFFFF;\n",
              "    }\n",
              "  </style>\n",
              "\n",
              "      <script>\n",
              "        const buttonEl =\n",
              "          document.querySelector('#df-667a7ab1-5964-45d2-ab73-627905bfd56f button.colab-df-convert');\n",
              "        buttonEl.style.display =\n",
              "          google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
              "\n",
              "        async function convertToInteractive(key) {\n",
              "          const element = document.querySelector('#df-667a7ab1-5964-45d2-ab73-627905bfd56f');\n",
              "          const dataTable =\n",
              "            await google.colab.kernel.invokeFunction('convertToInteractive',\n",
              "                                                     [key], {});\n",
              "          if (!dataTable) return;\n",
              "\n",
              "          const docLinkHtml = 'Like what you see? Visit the ' +\n",
              "            '<a target=\"_blank\" href=https://colab.research.google.com/notebooks/data_table.ipynb>data table notebook</a>'\n",
              "            + ' to learn more about interactive tables.';\n",
              "          element.innerHTML = '';\n",
              "          dataTable['output_type'] = 'display_data';\n",
              "          await google.colab.output.renderOutput(dataTable, element);\n",
              "          const docLink = document.createElement('div');\n",
              "          docLink.innerHTML = docLinkHtml;\n",
              "          element.appendChild(docLink);\n",
              "        }\n",
              "      </script>\n",
              "    </div>\n",
              "  </div>\n",
              "  "
            ]
          },
          "metadata": {},
          "execution_count": 31
        }
      ],
      "source": [
        "#Update Bert vocabylary\n",
        "import pandas as pd\n",
        "data = pd.read_csv(\"/content/relations.csv\")\n",
        "data\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UGRfDow9iNuc",
        "outputId": "f30ff685-b524-429a-8147-c707ddd6f074",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "rel\n",
              "other                   1639\n",
              "signed_by                865\n",
              "has_alias                471\n",
              "dated_as                 435\n",
              "has_collective_alias      91\n",
              "dtype: int64"
            ]
          },
          "metadata": {},
          "execution_count": 32
        }
      ],
      "source": [
        "data.value_counts('rel')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ApZ33Yx5iNuc",
        "outputId": "2165c9e8-7717-4d80-aac9-df0c9a8df464",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "label1\n",
              "PARTY                         1284\n",
              "DOC                           1267\n",
              "SIGNING_PERSON                 737\n",
              "ORG                             72\n",
              "ALIAS                           57\n",
              "ROLE                            41\n",
              "EFFDATE                         18\n",
              "SIGNING_TITLE                   11\n",
              "TITLE                            7\n",
              "NAME                             6\n",
              "PERMISSION_INDIRECT_OBJECT       1\n",
              "dtype: int64"
            ]
          },
          "metadata": {},
          "execution_count": 33
        }
      ],
      "source": [
        "data.value_counts('label1')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "19e-OdVYiNuc",
        "outputId": "1fc918cb-d94d-4b0d-d933-982e3cef3b7d",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "label2\n",
              "PARTY                882\n",
              "SIGNING_TITLE        724\n",
              "SIGNING_PERSON       498\n",
              "ALIAS                483\n",
              "ROLE                 389\n",
              "EFFDATE              376\n",
              "ORG                   63\n",
              "AGRDATE               52\n",
              "DOC                   19\n",
              "NAME                   7\n",
              "TITLE                  6\n",
              "FORMER_PARTY_NAME      1\n",
              "PERMISSION             1\n",
              "dtype: int64"
            ]
          },
          "metadata": {},
          "execution_count": 34
        }
      ],
      "source": [
        "data.value_counts('label2')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "WKj-K1dkgukV",
        "outputId": "26ca4223-92a6-43f6-fdda-553e9fdd3160"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "array(['signed_by', 'dated_as', 'has_alias', 'has_collective_alias',\n",
              "       'other'], dtype=object)"
            ]
          },
          "metadata": {},
          "execution_count": 35
        }
      ],
      "source": [
        "#get a list of valid relation names (less than 10 occurrences are probably wrong labels, or at least with a very low representation)\n",
        "valid_rel_labels = data['rel'].unique()\n",
        "valid_rel_labels"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "C3aEQUffhJdH"
      },
      "source": [
        "### Create the training and test datasets"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "df_train = data.sample(frac=0.9, random_state=1)\n",
        "df_test = data.drop(df_train.index)\n"
      ],
      "metadata": {
        "id": "IsyTzlzEHmAE"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "train_sentences, train_rel_label_ids, train_rel_arg_bindings, rel_labels = (\n",
        "    collect_data_from_pandas_dataset(df_train))\n",
        "\n",
        "test_sentences, test_rel_label_ids, test_rel_arg_bindings, _ = (\n",
        "    collect_data_from_pandas_dataset(df_test))"
      ],
      "metadata": {
        "id": "TdbXoebvH47P"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "train_sentences = train_sentences.values\n",
        "train_rel_label_ids = train_rel_label_ids.values\n",
        "train_rel_arg_bindings = train_rel_arg_bindings.values\n",
        "\n"
      ],
      "metadata": {
        "id": "uEAuuqdWH9xI"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "test_sentences = test_sentences.values\n",
        "test_rel_label_ids = test_rel_label_ids.values\n",
        "test_rel_arg_bindings = test_rel_arg_bindings.values"
      ],
      "metadata": {
        "id": "HnCyk2QXIBtR"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xa_cg2pxhSzL"
      },
      "source": [
        "### 7. Create the features from the datasets"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#create features\n",
        "train_features = make_features(\n",
        "    train_sentences, \n",
        "    list(zip(train_rel_label_ids, train_rel_arg_bindings)),\n",
        "    tokenizer, \n",
        "    is_test=False)\n",
        "\n",
        "if BertREConfig.REPLACE_ARG_PROB and BertREConfig.REPLICATE_DATASET:\n",
        "    train_features += make_features(\n",
        "        train_sentences, \n",
        "        list(zip(train_rel_label_ids, train_rel_arg_bindings)),\n",
        "        tokenizer, \n",
        "        is_test=True)\n",
        "        \n",
        "test_features = make_features(\n",
        "    test_sentences, \n",
        "    list(zip(test_rel_label_ids, test_rel_arg_bindings)),\n",
        "    tokenizer, \n",
        "    is_test=True)"
      ],
      "metadata": {
        "id": "WT6tGCCdg2y0"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aJLRWPMONuvj",
        "outputId": "a63b91bf-5033-4dd4-f738-6df5fd5bfc7b",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "2383 training examples\n",
            "266 test examples\n"
          ]
        }
      ],
      "source": [
        "print(\"{} training examples\".format(len(train_features)))\n",
        "print(\"{} test examples\".format(len(test_features)))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "N0EoU0UHNuvj"
      },
      "outputs": [],
      "source": [
        "train_features_part1 = train_features[:len(train_features) // 2]\n",
        "train_features_part2 = train_features[len(train_features) // 2:]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Qkb6vfMGNuvj",
        "outputId": "c5b7a012-9eb7-496b-81aa-d50c00bbb95c",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "1191 training part 1 examples\n",
            "1192 training part 2 examples\n",
            "266 test examples\n"
          ]
        }
      ],
      "source": [
        "print(\"{} training part 1 examples\".format(len(train_features_part1)))\n",
        "print(\"{} training part 2 examples\".format(len(train_features_part2)))\n",
        "print(\"{} test examples\".format(len(test_features)))"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from tensorflow.python.framework import ops"
      ],
      "metadata": {
        "id": "dRTFGZy7nxXt"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "rel_labels"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "moeDhfcDGc1k",
        "outputId": "363ab44d-900e-4afa-bba6-057e950e1767"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "['dated_as', 'has_alias', 'has_collective_alias', 'other', 'signed_by']"
            ]
          },
          "metadata": {},
          "execution_count": 58
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_dceO1UwiO8t"
      },
      "source": [
        "### Training the model\n",
        "Here is where all the fun happens! 🏄"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RN5N7vQWNuvj",
        "outputId": "1acd9dcd-dbd8-4d86-f992-adc6aef9f847",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "   Epoch      Batch         Loss     L_ACC   Arg_ACC       ACC    vL_ACC  vArg_ACC      vACC\n",
            "    1/5      148/148      0.4418     0.866     0.959     0.857     0.975     1.000     0.975\n",
            "    2/5      148/148      0.0478     0.987     1.000     0.987     0.975     1.000     0.975\n",
            "    3/5      148/148      0.0143     0.996     1.000     0.996     0.975     1.000     0.975\n",
            "    4/5      148/148      0.0049     0.999     1.000     0.999     0.970     1.000     0.970\n",
            "    5/5      148/148      0.0022     1.000     1.000     1.000     0.975     1.000     0.975\n"
          ]
        }
      ],
      "source": [
        "\n",
        "\n",
        "train_model(\n",
        "    \"CONTRACT_DOC_PARTIES_SPLIT\",\n",
        "    train_features=train_features, \n",
        "    test_features=test_features, \n",
        "    rel_labels=rel_labels,\n",
        "    num_arg_bindings=3)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9KVfT6p-oVG8"
      },
      "source": [
        "### Evaluating the model"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "tf.reset_default_graph()\n",
        "\n",
        "metrics = eval_metrics(\n",
        "    \"CONTRACT_DOC_PARTIES_SPLIT\", test_features, rel_labels, num_arg_bindings=0, exclude_rels=[])        \n",
        "print_metrics(metrics)"
      ],
      "metadata": {
        "id": "BHS8pX91XuEj",
        "outputId": "8058fd86-4abb-403f-9f9f-aa4a694712ec",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "\n",
            "Relation           Recall Precision        F1   Support\n",
            "\n",
            "dated_as            0.976     0.953     0.965        42\n",
            "has_alias           0.964     1.000     0.982        28\n",
            "has_collective_alias     1.000     1.000     1.000         1\n",
            "other               1.000     0.965     0.982        55\n",
            "signed_by           0.959     0.986     0.973        74\n",
            "\n",
            "Avg.                0.980     0.981     0.980\n",
            "\n",
            "Weighted Avg.       0.975     0.975     0.975\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "VGJQHgRJXuYL"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "UdEW-VCpbHNz"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Mnx9sUAO7OFr"
      },
      "source": [
        "# Train with all data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_3XfLpm77RzQ",
        "outputId": "d365df5b-0866-4967-e011-3c5bda38ead8",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "   Epoch      Batch         Loss     L_ACC   Arg_ACC       ACC    vL_ACC  vArg_ACC      vACC\n",
            "    1/5      165/165      0.3778     0.887     0.969     0.880       nan       nan       nan\n",
            "    2/5      165/165      0.0545     0.989     1.000     0.989       nan       nan       nan\n",
            "    3/5      165/165      0.0169     0.995     1.000     0.995       nan       nan       nan\n",
            "    4/5      165/165      0.0115     0.998     1.000     0.998       nan       nan       nan\n",
            "    5/5      165/165      0.0088     0.997     1.000     0.997       nan       nan       nan\n"
          ]
        }
      ],
      "source": [
        "tf.reset_default_graph()    \n",
        "\n",
        "train_model(\n",
        "    \"CONTRACT_DOC_PARTIES\", \n",
        "    train_features=train_features+test_features,\n",
        "    test_features=[], \n",
        "    rel_labels=rel_labels,\n",
        "    num_arg_bindings=3)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8nSMs3jzNuvj",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "4dee607a-9319-4a75-a5ea-e6e9fb255184"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "/content/trained/CONTRACT_DOC_PARTIES:\n",
            "total 1269052\n",
            "drwxr-xr-x 2 root root       4096 Feb 12 18:22 assets\n",
            "-rw-r--r-- 1 root root         67 Feb 12 18:22 checkpoint\n",
            "-rw-r--r-- 1 root root 1295219816 Feb 12 18:22 model.data-00000-of-00001\n",
            "-rw-r--r-- 1 root root      22852 Feb 12 18:22 model.index\n",
            "-rw-r--r-- 1 root root    4248536 Feb 12 18:22 model.meta\n",
            "\n",
            "/content/trained/CONTRACT_DOC_PARTIES_SPLIT:\n",
            "total 1269052\n",
            "drwxr-xr-x 2 root root       4096 Feb 12 18:15 assets\n",
            "-rw-r--r-- 1 root root         67 Feb 12 18:15 checkpoint\n",
            "-rw-r--r-- 1 root root 1295219816 Feb 12 18:15 model.data-00000-of-00001\n",
            "-rw-r--r-- 1 root root      22852 Feb 12 18:15 model.index\n",
            "-rw-r--r-- 1 root root    4248536 Feb 12 18:15 model.meta\n"
          ]
        }
      ],
      "source": [
        "!ls -l /content/trained/*"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "s7ef7-50obgT"
      },
      "source": [
        "### Finally saving it!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "code_folding": [],
        "id": "6yacQzYHNuvk",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "ed90908d-7f9d-4ec1-e9b1-6c0d8bf98999"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "108321797 trainbale parameters.\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "WARNING:tensorflow:From <ipython-input-67-0d4ad366df23>:51: simple_save (from tensorflow.python.saved_model.simple_save) is deprecated and will be removed in a future version.\n",
            "Instructions for updating:\n",
            "This API was designed for TensorFlow v1. See https://www.tensorflow.org/guide/migrate for instructions on how to migrate your code to TensorFlow v2.\n",
            "WARNING:tensorflow:From /usr/local/lib/python3.8/dist-packages/tensorflow/python/saved_model/signature_def_utils_impl.py:203: build_tensor_info (from tensorflow.python.saved_model.utils_impl) is deprecated and will be removed in a future version.\n",
            "Instructions for updating:\n",
            "This API was designed for TensorFlow v1. See https://www.tensorflow.org/guide/migrate for instructions on how to migrate your code to TensorFlow v2.\n"
          ]
        }
      ],
      "source": [
        "tf.reset_default_graph()\n",
        "\n",
        "export_model(\"CONTRACT_DOC_PARTIES\", is_trainable=False, num_arg_bindings=0)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AbNHyx1hNuvk",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "d6a01981-b3a5-442c-d4c8-096c8631f5ad"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "-rw-r--r-- 1 root root 1013294 Feb 12 18:22 /models/basebert_re/saved_model.pb\n",
            "\n",
            "/models/basebert_re/assets:\n",
            "total 216\n",
            "-rw-r--r-- 1 root root     55 Feb 12 18:22 categories.txt\n",
            "-rw-r--r-- 1 root root 213422 Feb 12 18:22 vocab.txt\n",
            "\n",
            "/models/basebert_re/variables:\n",
            "total 423148\n",
            "-rw-r--r-- 1 root root 433287188 Feb 12 18:22 variables.data-00000-of-00001\n",
            "-rw-r--r-- 1 root root      8306 Feb 12 18:22 variables.index\n"
          ]
        }
      ],
      "source": [
        "!ls -l  {BertREConfig.EXPORT_PATH}/*"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OOBOjF1giNue",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "f56af584-2c42-43e3-d243-126ca7662a45"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Reading package lists... Done\n",
            "Building dependency tree       \n",
            "Reading state information... Done\n",
            "zip is already the newest version (3.0-11build1).\n",
            "unzip is already the newest version (6.0-25ubuntu1.1).\n",
            "The following package was automatically installed and is no longer required:\n",
            "  libnvidia-common-510\n",
            "Use 'sudo apt autoremove' to remove it.\n",
            "0 upgraded, 0 newly installed, 0 to remove and 21 not upgraded.\n"
          ]
        }
      ],
      "source": [
        "!sudo apt-get -y install zip unzip"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JteBYaHueUMF",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "eb2383fc-f88b-4277-e134-b13d70e48748"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "  adding: content/models/basebert_re/ (stored 0%)\n",
            "  adding: content/models/basebert_re/saved_model.pb (deflated 92%)\n",
            "  adding: content/models/basebert_re/variables/ (stored 0%)\n",
            "  adding: content/models/basebert_re/variables/variables.index (deflated 68%)\n",
            "  adding: content/models/basebert_re/variables/variables.data-00000-of-00001 (deflated 7%)\n",
            "  adding: content/models/basebert_re/assets/ (stored 0%)\n",
            "  adding: content/models/basebert_re/assets/categories.txt (deflated 18%)\n",
            "  adding: content/models/basebert_re/assets/vocab.txt (deflated 49%)\n"
          ]
        }
      ],
      "source": [
        "!zip -r redl.zip /content/models/basebert_re"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oYSUxcZtiNue"
      },
      "source": [
        "# We test in SPARK NLP\n",
        "\n",
        "Let's test our model with SparkNLP"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cokHOoyViNue"
      },
      "outputs": [],
      "source": [
        "import pandas as pd\n",
        "\n",
        "def get_relations_df (results, col='relations'):\n",
        "  rel_pairs=[]\n",
        "  for rel in results[0][col]:\n",
        "      rel_pairs.append((\n",
        "          rel.result, \n",
        "          rel.metadata['entity1'], \n",
        "          rel.metadata['entity1_begin'],\n",
        "          rel.metadata['entity1_end'],\n",
        "          rel.metadata['chunk1'], \n",
        "          rel.metadata['entity2'],\n",
        "          rel.metadata['entity2_begin'],\n",
        "          rel.metadata['entity2_end'],\n",
        "          rel.metadata['chunk2'], \n",
        "          rel.metadata['confidence']\n",
        "      ))\n",
        "\n",
        "  rel_df = pd.DataFrame(rel_pairs, columns=['relation','entity1','entity1_begin','entity1_end','chunk1','entity2','entity2_begin','entity2_end','chunk2', 'confidence'])\n",
        "\n",
        "  return rel_df"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_W-OfQw3jWtJ"
      },
      "source": [
        "Here, we import our model to SparkNLP"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wdt5RR3OiNue"
      },
      "outputs": [],
      "source": [
        "re = legal.RelationExtractionDLModel().loadSavedModel('/content/models/basebert_re', spark)\n",
        "re.write().overwrite().save('legre_contract_doc_parties')"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "text='''\n",
        "This INTELLECTUAL PROPERTY AGREEMENT (this \"Agreement\"), dated as of December 31, 2018 (the \"Effective Date\") is entered into by and between Armstrong Flooring, Inc., a Delaware corporation (\"Seller\") and AFI Licensing LLC, a Delaware limited liability company (\"Licensing\" and together with Seller, \"Arizona\") and AHF Holding, Inc. (formerly known as Tarzan HoldCo, Inc.), a Delaware corporation (\"Buyer\") and Armstrong Hardwood Flooring Company, a Tennessee corporation (the \"Company\" and together with Buyer the \"Buyer Entities\") (each of Arizona on the one hand and the Buyer Entities on the other hand, a \"Party\" and collectively, the \"Parties\").\n",
        "'''"
      ],
      "metadata": {
        "id": "wlnLOR7mRbN3"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8Ah9Rn47jfhl"
      },
      "source": [
        "Now before getting relations, we have to extract entities from the given text. For this, we will use `legner_contract_doc_parties_lg` NER model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qSj5PJOxiNue",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "0a06cb59-4a53-4518-9edb-feb3c87c4e91"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "roberta_embeddings_legal_roberta_base download started this may take some time.\n",
            "Approximate size to download 447.2 MB\n",
            "[OK!]\n",
            "legner_contract_doc_parties_lg download started this may take some time.\n",
            "[OK!]\n"
          ]
        }
      ],
      "source": [
        "document_assembler = nlp.DocumentAssembler()\\\n",
        "  .setInputCol(\"text\")\\\n",
        "  .setOutputCol(\"document\")\n",
        "\n",
        "tokenizer = nlp.Tokenizer()\\\n",
        "    .setInputCols(\"document\")\\\n",
        "    .setOutputCol(\"token\")\n",
        "\n",
        "embeddings = nlp.RoBertaEmbeddings.pretrained(\"roberta_embeddings_legal_roberta_base\", \"en\") \\\n",
        "    .setInputCols(\"document\", \"token\") \\\n",
        "    .setOutputCol(\"embeddings\")\\\n",
        "    .setMaxSentenceLength(512)\n",
        "\n",
        "ner_model = legal.NerModel.pretrained('legner_contract_doc_parties_lg', 'en', 'legal/models')\\\n",
        "    .setInputCols([\"sentence\", \"token\", \"embeddings\"])\\\n",
        "    .setOutputCol(\"ner\")\n",
        "\n",
        "ner_converter = nlp.NerConverter()\\\n",
        "    .setInputCols([\"document\",\"token\",\"ner\"])\\\n",
        "    .setOutputCol(\"ner_chunk\")\n",
        "\n",
        "# We use the load function to run our trained model.\n",
        "reDL = legal.RelationExtractionDLModel().load('legre_contract_doc_parties')\\\n",
        "    .setPredictionThreshold(0.5)\\\n",
        "    .setInputCols([\"ner_chunk\", \"document\"])\\\n",
        "    .setOutputCol(\"relations\")\n",
        "\n",
        "nlpPipeline = nlp.Pipeline(stages=[\n",
        "    document_assembler,\n",
        "    tokenizer,\n",
        "    embeddings,\n",
        "    ner_model,\n",
        "    ner_converter,\n",
        "    reDL\n",
        "    ])\n",
        "\n",
        "data = spark.createDataFrame([[text]]).toDF(\"text\")\n",
        "\n",
        "model = nlpPipeline.fit(data)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tDZ2tLeviNue"
      },
      "outputs": [],
      "source": [
        "light_model = nlp.LightPipeline(model)\n",
        "\n",
        "\n",
        "results = light_model.fullAnnotate(text)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RCkP3vFgiNue",
        "outputId": "eaab7575-c28f-44bb-f97a-07fca29a8ac1",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        }
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "                relation  entity1 entity1_begin entity1_end  \\\n",
              "0               dated_as      DOC             6          36   \n",
              "1              signed_by      DOC             6          36   \n",
              "2              signed_by      DOC             6          36   \n",
              "3              signed_by      DOC             6          36   \n",
              "4              signed_by      DOC             6          36   \n",
              "5              signed_by      DOC             6          36   \n",
              "6              signed_by      DOC             6          36   \n",
              "8              signed_by      DOC             6          36   \n",
              "9              signed_by      DOC             6          36   \n",
              "10             signed_by      DOC             6          36   \n",
              "11             signed_by      DOC             6          36   \n",
              "12             signed_by      DOC             6          36   \n",
              "13             signed_by      DOC             6          36   \n",
              "14             signed_by  EFFDATE            70          86   \n",
              "15              dated_as  EFFDATE            70          86   \n",
              "16             signed_by  EFFDATE            70          86   \n",
              "17              dated_as  EFFDATE            70          86   \n",
              "18              dated_as  EFFDATE            70          86   \n",
              "19             signed_by  EFFDATE            70          86   \n",
              "21             signed_by  EFFDATE            70          86   \n",
              "22              dated_as  EFFDATE            70          86   \n",
              "23              dated_as  EFFDATE            70          86   \n",
              "24              dated_as  EFFDATE            70          86   \n",
              "25              dated_as  EFFDATE            70          86   \n",
              "26              dated_as  EFFDATE            70          86   \n",
              "39             signed_by    ALIAS           193         198   \n",
              "41  has_collective_alias    ALIAS           193         198   \n",
              "43  has_collective_alias    ALIAS           193         198   \n",
              "44  has_collective_alias    ALIAS           193         198   \n",
              "45  has_collective_alias    ALIAS           193         198   \n",
              "46  has_collective_alias    ALIAS           193         198   \n",
              "47  has_collective_alias    ALIAS           193         198   \n",
              "48  has_collective_alias    ALIAS           193         198   \n",
              "59  has_collective_alias    ALIAS           264         272   \n",
              "60             signed_by    ALIAS           264         272   \n",
              "62  has_collective_alias    ALIAS           264         272   \n",
              "63  has_collective_alias    ALIAS           264         272   \n",
              "64  has_collective_alias    ALIAS           264         272   \n",
              "65  has_collective_alias    ALIAS           264         272   \n",
              "66  has_collective_alias    ALIAS           264         272   \n",
              "67  has_collective_alias    ALIAS           264         272   \n",
              "68             signed_by    PARTY           293         298   \n",
              "70  has_collective_alias    PARTY           293         298   \n",
              "71  has_collective_alias    PARTY           293         298   \n",
              "72  has_collective_alias    PARTY           293         298   \n",
              "73  has_collective_alias    PARTY           293         298   \n",
              "74  has_collective_alias    PARTY           293         298   \n",
              "82  has_collective_alias    ALIAS           400         404   \n",
              "83  has_collective_alias    ALIAS           400         404   \n",
              "84  has_collective_alias    ALIAS           400         404   \n",
              "85  has_collective_alias    ALIAS           400         404   \n",
              "86  has_collective_alias    ALIAS           400         404   \n",
              "\n",
              "                             chunk1  entity2 entity2_begin entity2_end  \\\n",
              "0   INTELLECTUAL PROPERTY AGREEMENT  EFFDATE            70          86   \n",
              "1   INTELLECTUAL PROPERTY AGREEMENT    PARTY           142         164   \n",
              "2   INTELLECTUAL PROPERTY AGREEMENT    ALIAS           193         198   \n",
              "3   INTELLECTUAL PROPERTY AGREEMENT    PARTY           206         222   \n",
              "4   INTELLECTUAL PROPERTY AGREEMENT    ALIAS           264         272   \n",
              "5   INTELLECTUAL PROPERTY AGREEMENT    PARTY           293         298   \n",
              "6   INTELLECTUAL PROPERTY AGREEMENT    PARTY           316         331   \n",
              "8   INTELLECTUAL PROPERTY AGREEMENT    PARTY           412         446   \n",
              "9   INTELLECTUAL PROPERTY AGREEMENT    ALIAS           479         485   \n",
              "10  INTELLECTUAL PROPERTY AGREEMENT    ALIAS           517         530   \n",
              "11  INTELLECTUAL PROPERTY AGREEMENT    PARTY           575         588   \n",
              "12  INTELLECTUAL PROPERTY AGREEMENT    ALIAS           612         616   \n",
              "13  INTELLECTUAL PROPERTY AGREEMENT    ALIAS           642         648   \n",
              "14                December 31, 2018    PARTY           142         164   \n",
              "15                December 31, 2018    ALIAS           193         198   \n",
              "16                December 31, 2018    PARTY           206         222   \n",
              "17                December 31, 2018    ALIAS           264         272   \n",
              "18                December 31, 2018    PARTY           293         298   \n",
              "19                December 31, 2018    PARTY           316         331   \n",
              "21                December 31, 2018    PARTY           412         446   \n",
              "22                December 31, 2018    ALIAS           479         485   \n",
              "23                December 31, 2018    ALIAS           517         530   \n",
              "24                December 31, 2018    PARTY           575         588   \n",
              "25                December 31, 2018    ALIAS           612         616   \n",
              "26                December 31, 2018    ALIAS           642         648   \n",
              "39                           Seller    PARTY           206         222   \n",
              "41                           Seller    PARTY           293         298   \n",
              "43                           Seller    PARTY           412         446   \n",
              "44                           Seller    ALIAS           479         485   \n",
              "45                           Seller    ALIAS           517         530   \n",
              "46                           Seller    PARTY           575         588   \n",
              "47                           Seller    ALIAS           612         616   \n",
              "48                           Seller    ALIAS           642         648   \n",
              "59                        Licensing    PARTY           293         298   \n",
              "60                        Licensing    PARTY           316         331   \n",
              "62                        Licensing    PARTY           412         446   \n",
              "63                        Licensing    ALIAS           479         485   \n",
              "64                        Licensing    ALIAS           517         530   \n",
              "65                        Licensing    PARTY           575         588   \n",
              "66                        Licensing    ALIAS           612         616   \n",
              "67                        Licensing    ALIAS           642         648   \n",
              "68                           Seller    PARTY           316         331   \n",
              "70                           Seller    ALIAS           479         485   \n",
              "71                           Seller    ALIAS           517         530   \n",
              "72                           Seller    PARTY           575         588   \n",
              "73                           Seller    ALIAS           612         616   \n",
              "74                           Seller    ALIAS           642         648   \n",
              "82                            Buyer    ALIAS           479         485   \n",
              "83                            Buyer    ALIAS           517         530   \n",
              "84                            Buyer    PARTY           575         588   \n",
              "85                            Buyer    ALIAS           612         616   \n",
              "86                            Buyer    ALIAS           642         648   \n",
              "\n",
              "                                 chunk2  confidence  \n",
              "0                     December 31, 2018   0.9998894  \n",
              "1               Armstrong Flooring, Inc  0.99805707  \n",
              "2                                Seller  0.54615384  \n",
              "3                     AFI Licensing LLC   0.9963246  \n",
              "4                             Licensing   0.5839682  \n",
              "5                                Seller   0.9191695  \n",
              "6                      AHF Holding, Inc  0.99637455  \n",
              "8   Armstrong Hardwood Flooring Company   0.9782145  \n",
              "9                               Company  0.93827885  \n",
              "10                       Buyer Entities  0.93827885  \n",
              "11                       Buyer Entities  0.93827885  \n",
              "12                                Party  0.93827885  \n",
              "13                              Parties  0.93827885  \n",
              "14              Armstrong Flooring, Inc  0.72774583  \n",
              "15                               Seller   0.6977681  \n",
              "16                    AFI Licensing LLC  0.86870426  \n",
              "17                            Licensing  0.73973966  \n",
              "18                               Seller  0.70528424  \n",
              "19                     AHF Holding, Inc    0.904696  \n",
              "21  Armstrong Hardwood Flooring Company  0.62537766  \n",
              "22                              Company   0.9445322  \n",
              "23                       Buyer Entities   0.9445322  \n",
              "24                       Buyer Entities   0.9445322  \n",
              "25                                Party   0.9445322  \n",
              "26                              Parties   0.9445322  \n",
              "39                    AFI Licensing LLC  0.53952813  \n",
              "41                               Seller  0.86056995  \n",
              "43  Armstrong Hardwood Flooring Company  0.61624825  \n",
              "44                              Company   0.8246864  \n",
              "45                       Buyer Entities   0.8246864  \n",
              "46                       Buyer Entities   0.8246864  \n",
              "47                                Party   0.8246864  \n",
              "48                              Parties   0.8246864  \n",
              "59                               Seller   0.9327616  \n",
              "60                     AHF Holding, Inc  0.50352174  \n",
              "62  Armstrong Hardwood Flooring Company  0.75796545  \n",
              "63                              Company  0.85309875  \n",
              "64                       Buyer Entities  0.85309875  \n",
              "65                       Buyer Entities  0.85309875  \n",
              "66                                Party  0.85309875  \n",
              "67                              Parties  0.85309875  \n",
              "68                     AHF Holding, Inc   0.5209184  \n",
              "70                              Company   0.5416529  \n",
              "71                       Buyer Entities   0.5416529  \n",
              "72                       Buyer Entities   0.5416529  \n",
              "73                                Party   0.5416529  \n",
              "74                              Parties   0.5416529  \n",
              "82                              Company  0.60943455  \n",
              "83                       Buyer Entities  0.60943455  \n",
              "84                       Buyer Entities  0.60943455  \n",
              "85                                Party  0.60943455  \n",
              "86                              Parties  0.60943455  "
            ],
            "text/html": [
              "\n",
              "  <div id=\"df-f6d50ef3-a2ba-4632-91dc-b5dd7423352d\">\n",
              "    <div class=\"colab-df-container\">\n",
              "      <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>relation</th>\n",
              "      <th>entity1</th>\n",
              "      <th>entity1_begin</th>\n",
              "      <th>entity1_end</th>\n",
              "      <th>chunk1</th>\n",
              "      <th>entity2</th>\n",
              "      <th>entity2_begin</th>\n",
              "      <th>entity2_end</th>\n",
              "      <th>chunk2</th>\n",
              "      <th>confidence</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>0</th>\n",
              "      <td>dated_as</td>\n",
              "      <td>DOC</td>\n",
              "      <td>6</td>\n",
              "      <td>36</td>\n",
              "      <td>INTELLECTUAL PROPERTY AGREEMENT</td>\n",
              "      <td>EFFDATE</td>\n",
              "      <td>70</td>\n",
              "      <td>86</td>\n",
              "      <td>December 31, 2018</td>\n",
              "      <td>0.9998894</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>1</th>\n",
              "      <td>signed_by</td>\n",
              "      <td>DOC</td>\n",
              "      <td>6</td>\n",
              "      <td>36</td>\n",
              "      <td>INTELLECTUAL PROPERTY AGREEMENT</td>\n",
              "      <td>PARTY</td>\n",
              "      <td>142</td>\n",
              "      <td>164</td>\n",
              "      <td>Armstrong Flooring, Inc</td>\n",
              "      <td>0.99805707</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>2</th>\n",
              "      <td>signed_by</td>\n",
              "      <td>DOC</td>\n",
              "      <td>6</td>\n",
              "      <td>36</td>\n",
              "      <td>INTELLECTUAL PROPERTY AGREEMENT</td>\n",
              "      <td>ALIAS</td>\n",
              "      <td>193</td>\n",
              "      <td>198</td>\n",
              "      <td>Seller</td>\n",
              "      <td>0.54615384</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>3</th>\n",
              "      <td>signed_by</td>\n",
              "      <td>DOC</td>\n",
              "      <td>6</td>\n",
              "      <td>36</td>\n",
              "      <td>INTELLECTUAL PROPERTY AGREEMENT</td>\n",
              "      <td>PARTY</td>\n",
              "      <td>206</td>\n",
              "      <td>222</td>\n",
              "      <td>AFI Licensing LLC</td>\n",
              "      <td>0.9963246</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>4</th>\n",
              "      <td>signed_by</td>\n",
              "      <td>DOC</td>\n",
              "      <td>6</td>\n",
              "      <td>36</td>\n",
              "      <td>INTELLECTUAL PROPERTY AGREEMENT</td>\n",
              "      <td>ALIAS</td>\n",
              "      <td>264</td>\n",
              "      <td>272</td>\n",
              "      <td>Licensing</td>\n",
              "      <td>0.5839682</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>5</th>\n",
              "      <td>signed_by</td>\n",
              "      <td>DOC</td>\n",
              "      <td>6</td>\n",
              "      <td>36</td>\n",
              "      <td>INTELLECTUAL PROPERTY AGREEMENT</td>\n",
              "      <td>PARTY</td>\n",
              "      <td>293</td>\n",
              "      <td>298</td>\n",
              "      <td>Seller</td>\n",
              "      <td>0.9191695</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>6</th>\n",
              "      <td>signed_by</td>\n",
              "      <td>DOC</td>\n",
              "      <td>6</td>\n",
              "      <td>36</td>\n",
              "      <td>INTELLECTUAL PROPERTY AGREEMENT</td>\n",
              "      <td>PARTY</td>\n",
              "      <td>316</td>\n",
              "      <td>331</td>\n",
              "      <td>AHF Holding, Inc</td>\n",
              "      <td>0.99637455</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>8</th>\n",
              "      <td>signed_by</td>\n",
              "      <td>DOC</td>\n",
              "      <td>6</td>\n",
              "      <td>36</td>\n",
              "      <td>INTELLECTUAL PROPERTY AGREEMENT</td>\n",
              "      <td>PARTY</td>\n",
              "      <td>412</td>\n",
              "      <td>446</td>\n",
              "      <td>Armstrong Hardwood Flooring Company</td>\n",
              "      <td>0.9782145</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>9</th>\n",
              "      <td>signed_by</td>\n",
              "      <td>DOC</td>\n",
              "      <td>6</td>\n",
              "      <td>36</td>\n",
              "      <td>INTELLECTUAL PROPERTY AGREEMENT</td>\n",
              "      <td>ALIAS</td>\n",
              "      <td>479</td>\n",
              "      <td>485</td>\n",
              "      <td>Company</td>\n",
              "      <td>0.93827885</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>10</th>\n",
              "      <td>signed_by</td>\n",
              "      <td>DOC</td>\n",
              "      <td>6</td>\n",
              "      <td>36</td>\n",
              "      <td>INTELLECTUAL PROPERTY AGREEMENT</td>\n",
              "      <td>ALIAS</td>\n",
              "      <td>517</td>\n",
              "      <td>530</td>\n",
              "      <td>Buyer Entities</td>\n",
              "      <td>0.93827885</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>11</th>\n",
              "      <td>signed_by</td>\n",
              "      <td>DOC</td>\n",
              "      <td>6</td>\n",
              "      <td>36</td>\n",
              "      <td>INTELLECTUAL PROPERTY AGREEMENT</td>\n",
              "      <td>PARTY</td>\n",
              "      <td>575</td>\n",
              "      <td>588</td>\n",
              "      <td>Buyer Entities</td>\n",
              "      <td>0.93827885</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>12</th>\n",
              "      <td>signed_by</td>\n",
              "      <td>DOC</td>\n",
              "      <td>6</td>\n",
              "      <td>36</td>\n",
              "      <td>INTELLECTUAL PROPERTY AGREEMENT</td>\n",
              "      <td>ALIAS</td>\n",
              "      <td>612</td>\n",
              "      <td>616</td>\n",
              "      <td>Party</td>\n",
              "      <td>0.93827885</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>13</th>\n",
              "      <td>signed_by</td>\n",
              "      <td>DOC</td>\n",
              "      <td>6</td>\n",
              "      <td>36</td>\n",
              "      <td>INTELLECTUAL PROPERTY AGREEMENT</td>\n",
              "      <td>ALIAS</td>\n",
              "      <td>642</td>\n",
              "      <td>648</td>\n",
              "      <td>Parties</td>\n",
              "      <td>0.93827885</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>14</th>\n",
              "      <td>signed_by</td>\n",
              "      <td>EFFDATE</td>\n",
              "      <td>70</td>\n",
              "      <td>86</td>\n",
              "      <td>December 31, 2018</td>\n",
              "      <td>PARTY</td>\n",
              "      <td>142</td>\n",
              "      <td>164</td>\n",
              "      <td>Armstrong Flooring, Inc</td>\n",
              "      <td>0.72774583</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>15</th>\n",
              "      <td>dated_as</td>\n",
              "      <td>EFFDATE</td>\n",
              "      <td>70</td>\n",
              "      <td>86</td>\n",
              "      <td>December 31, 2018</td>\n",
              "      <td>ALIAS</td>\n",
              "      <td>193</td>\n",
              "      <td>198</td>\n",
              "      <td>Seller</td>\n",
              "      <td>0.6977681</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>16</th>\n",
              "      <td>signed_by</td>\n",
              "      <td>EFFDATE</td>\n",
              "      <td>70</td>\n",
              "      <td>86</td>\n",
              "      <td>December 31, 2018</td>\n",
              "      <td>PARTY</td>\n",
              "      <td>206</td>\n",
              "      <td>222</td>\n",
              "      <td>AFI Licensing LLC</td>\n",
              "      <td>0.86870426</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>17</th>\n",
              "      <td>dated_as</td>\n",
              "      <td>EFFDATE</td>\n",
              "      <td>70</td>\n",
              "      <td>86</td>\n",
              "      <td>December 31, 2018</td>\n",
              "      <td>ALIAS</td>\n",
              "      <td>264</td>\n",
              "      <td>272</td>\n",
              "      <td>Licensing</td>\n",
              "      <td>0.73973966</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>18</th>\n",
              "      <td>dated_as</td>\n",
              "      <td>EFFDATE</td>\n",
              "      <td>70</td>\n",
              "      <td>86</td>\n",
              "      <td>December 31, 2018</td>\n",
              "      <td>PARTY</td>\n",
              "      <td>293</td>\n",
              "      <td>298</td>\n",
              "      <td>Seller</td>\n",
              "      <td>0.70528424</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>19</th>\n",
              "      <td>signed_by</td>\n",
              "      <td>EFFDATE</td>\n",
              "      <td>70</td>\n",
              "      <td>86</td>\n",
              "      <td>December 31, 2018</td>\n",
              "      <td>PARTY</td>\n",
              "      <td>316</td>\n",
              "      <td>331</td>\n",
              "      <td>AHF Holding, Inc</td>\n",
              "      <td>0.904696</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>21</th>\n",
              "      <td>signed_by</td>\n",
              "      <td>EFFDATE</td>\n",
              "      <td>70</td>\n",
              "      <td>86</td>\n",
              "      <td>December 31, 2018</td>\n",
              "      <td>PARTY</td>\n",
              "      <td>412</td>\n",
              "      <td>446</td>\n",
              "      <td>Armstrong Hardwood Flooring Company</td>\n",
              "      <td>0.62537766</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>22</th>\n",
              "      <td>dated_as</td>\n",
              "      <td>EFFDATE</td>\n",
              "      <td>70</td>\n",
              "      <td>86</td>\n",
              "      <td>December 31, 2018</td>\n",
              "      <td>ALIAS</td>\n",
              "      <td>479</td>\n",
              "      <td>485</td>\n",
              "      <td>Company</td>\n",
              "      <td>0.9445322</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>23</th>\n",
              "      <td>dated_as</td>\n",
              "      <td>EFFDATE</td>\n",
              "      <td>70</td>\n",
              "      <td>86</td>\n",
              "      <td>December 31, 2018</td>\n",
              "      <td>ALIAS</td>\n",
              "      <td>517</td>\n",
              "      <td>530</td>\n",
              "      <td>Buyer Entities</td>\n",
              "      <td>0.9445322</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>24</th>\n",
              "      <td>dated_as</td>\n",
              "      <td>EFFDATE</td>\n",
              "      <td>70</td>\n",
              "      <td>86</td>\n",
              "      <td>December 31, 2018</td>\n",
              "      <td>PARTY</td>\n",
              "      <td>575</td>\n",
              "      <td>588</td>\n",
              "      <td>Buyer Entities</td>\n",
              "      <td>0.9445322</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>25</th>\n",
              "      <td>dated_as</td>\n",
              "      <td>EFFDATE</td>\n",
              "      <td>70</td>\n",
              "      <td>86</td>\n",
              "      <td>December 31, 2018</td>\n",
              "      <td>ALIAS</td>\n",
              "      <td>612</td>\n",
              "      <td>616</td>\n",
              "      <td>Party</td>\n",
              "      <td>0.9445322</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>26</th>\n",
              "      <td>dated_as</td>\n",
              "      <td>EFFDATE</td>\n",
              "      <td>70</td>\n",
              "      <td>86</td>\n",
              "      <td>December 31, 2018</td>\n",
              "      <td>ALIAS</td>\n",
              "      <td>642</td>\n",
              "      <td>648</td>\n",
              "      <td>Parties</td>\n",
              "      <td>0.9445322</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>39</th>\n",
              "      <td>signed_by</td>\n",
              "      <td>ALIAS</td>\n",
              "      <td>193</td>\n",
              "      <td>198</td>\n",
              "      <td>Seller</td>\n",
              "      <td>PARTY</td>\n",
              "      <td>206</td>\n",
              "      <td>222</td>\n",
              "      <td>AFI Licensing LLC</td>\n",
              "      <td>0.53952813</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>41</th>\n",
              "      <td>has_collective_alias</td>\n",
              "      <td>ALIAS</td>\n",
              "      <td>193</td>\n",
              "      <td>198</td>\n",
              "      <td>Seller</td>\n",
              "      <td>PARTY</td>\n",
              "      <td>293</td>\n",
              "      <td>298</td>\n",
              "      <td>Seller</td>\n",
              "      <td>0.86056995</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>43</th>\n",
              "      <td>has_collective_alias</td>\n",
              "      <td>ALIAS</td>\n",
              "      <td>193</td>\n",
              "      <td>198</td>\n",
              "      <td>Seller</td>\n",
              "      <td>PARTY</td>\n",
              "      <td>412</td>\n",
              "      <td>446</td>\n",
              "      <td>Armstrong Hardwood Flooring Company</td>\n",
              "      <td>0.61624825</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>44</th>\n",
              "      <td>has_collective_alias</td>\n",
              "      <td>ALIAS</td>\n",
              "      <td>193</td>\n",
              "      <td>198</td>\n",
              "      <td>Seller</td>\n",
              "      <td>ALIAS</td>\n",
              "      <td>479</td>\n",
              "      <td>485</td>\n",
              "      <td>Company</td>\n",
              "      <td>0.8246864</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>45</th>\n",
              "      <td>has_collective_alias</td>\n",
              "      <td>ALIAS</td>\n",
              "      <td>193</td>\n",
              "      <td>198</td>\n",
              "      <td>Seller</td>\n",
              "      <td>ALIAS</td>\n",
              "      <td>517</td>\n",
              "      <td>530</td>\n",
              "      <td>Buyer Entities</td>\n",
              "      <td>0.8246864</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>46</th>\n",
              "      <td>has_collective_alias</td>\n",
              "      <td>ALIAS</td>\n",
              "      <td>193</td>\n",
              "      <td>198</td>\n",
              "      <td>Seller</td>\n",
              "      <td>PARTY</td>\n",
              "      <td>575</td>\n",
              "      <td>588</td>\n",
              "      <td>Buyer Entities</td>\n",
              "      <td>0.8246864</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>47</th>\n",
              "      <td>has_collective_alias</td>\n",
              "      <td>ALIAS</td>\n",
              "      <td>193</td>\n",
              "      <td>198</td>\n",
              "      <td>Seller</td>\n",
              "      <td>ALIAS</td>\n",
              "      <td>612</td>\n",
              "      <td>616</td>\n",
              "      <td>Party</td>\n",
              "      <td>0.8246864</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>48</th>\n",
              "      <td>has_collective_alias</td>\n",
              "      <td>ALIAS</td>\n",
              "      <td>193</td>\n",
              "      <td>198</td>\n",
              "      <td>Seller</td>\n",
              "      <td>ALIAS</td>\n",
              "      <td>642</td>\n",
              "      <td>648</td>\n",
              "      <td>Parties</td>\n",
              "      <td>0.8246864</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>59</th>\n",
              "      <td>has_collective_alias</td>\n",
              "      <td>ALIAS</td>\n",
              "      <td>264</td>\n",
              "      <td>272</td>\n",
              "      <td>Licensing</td>\n",
              "      <td>PARTY</td>\n",
              "      <td>293</td>\n",
              "      <td>298</td>\n",
              "      <td>Seller</td>\n",
              "      <td>0.9327616</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>60</th>\n",
              "      <td>signed_by</td>\n",
              "      <td>ALIAS</td>\n",
              "      <td>264</td>\n",
              "      <td>272</td>\n",
              "      <td>Licensing</td>\n",
              "      <td>PARTY</td>\n",
              "      <td>316</td>\n",
              "      <td>331</td>\n",
              "      <td>AHF Holding, Inc</td>\n",
              "      <td>0.50352174</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>62</th>\n",
              "      <td>has_collective_alias</td>\n",
              "      <td>ALIAS</td>\n",
              "      <td>264</td>\n",
              "      <td>272</td>\n",
              "      <td>Licensing</td>\n",
              "      <td>PARTY</td>\n",
              "      <td>412</td>\n",
              "      <td>446</td>\n",
              "      <td>Armstrong Hardwood Flooring Company</td>\n",
              "      <td>0.75796545</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>63</th>\n",
              "      <td>has_collective_alias</td>\n",
              "      <td>ALIAS</td>\n",
              "      <td>264</td>\n",
              "      <td>272</td>\n",
              "      <td>Licensing</td>\n",
              "      <td>ALIAS</td>\n",
              "      <td>479</td>\n",
              "      <td>485</td>\n",
              "      <td>Company</td>\n",
              "      <td>0.85309875</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>64</th>\n",
              "      <td>has_collective_alias</td>\n",
              "      <td>ALIAS</td>\n",
              "      <td>264</td>\n",
              "      <td>272</td>\n",
              "      <td>Licensing</td>\n",
              "      <td>ALIAS</td>\n",
              "      <td>517</td>\n",
              "      <td>530</td>\n",
              "      <td>Buyer Entities</td>\n",
              "      <td>0.85309875</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>65</th>\n",
              "      <td>has_collective_alias</td>\n",
              "      <td>ALIAS</td>\n",
              "      <td>264</td>\n",
              "      <td>272</td>\n",
              "      <td>Licensing</td>\n",
              "      <td>PARTY</td>\n",
              "      <td>575</td>\n",
              "      <td>588</td>\n",
              "      <td>Buyer Entities</td>\n",
              "      <td>0.85309875</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>66</th>\n",
              "      <td>has_collective_alias</td>\n",
              "      <td>ALIAS</td>\n",
              "      <td>264</td>\n",
              "      <td>272</td>\n",
              "      <td>Licensing</td>\n",
              "      <td>ALIAS</td>\n",
              "      <td>612</td>\n",
              "      <td>616</td>\n",
              "      <td>Party</td>\n",
              "      <td>0.85309875</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>67</th>\n",
              "      <td>has_collective_alias</td>\n",
              "      <td>ALIAS</td>\n",
              "      <td>264</td>\n",
              "      <td>272</td>\n",
              "      <td>Licensing</td>\n",
              "      <td>ALIAS</td>\n",
              "      <td>642</td>\n",
              "      <td>648</td>\n",
              "      <td>Parties</td>\n",
              "      <td>0.85309875</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>68</th>\n",
              "      <td>signed_by</td>\n",
              "      <td>PARTY</td>\n",
              "      <td>293</td>\n",
              "      <td>298</td>\n",
              "      <td>Seller</td>\n",
              "      <td>PARTY</td>\n",
              "      <td>316</td>\n",
              "      <td>331</td>\n",
              "      <td>AHF Holding, Inc</td>\n",
              "      <td>0.5209184</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>70</th>\n",
              "      <td>has_collective_alias</td>\n",
              "      <td>PARTY</td>\n",
              "      <td>293</td>\n",
              "      <td>298</td>\n",
              "      <td>Seller</td>\n",
              "      <td>ALIAS</td>\n",
              "      <td>479</td>\n",
              "      <td>485</td>\n",
              "      <td>Company</td>\n",
              "      <td>0.5416529</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>71</th>\n",
              "      <td>has_collective_alias</td>\n",
              "      <td>PARTY</td>\n",
              "      <td>293</td>\n",
              "      <td>298</td>\n",
              "      <td>Seller</td>\n",
              "      <td>ALIAS</td>\n",
              "      <td>517</td>\n",
              "      <td>530</td>\n",
              "      <td>Buyer Entities</td>\n",
              "      <td>0.5416529</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>72</th>\n",
              "      <td>has_collective_alias</td>\n",
              "      <td>PARTY</td>\n",
              "      <td>293</td>\n",
              "      <td>298</td>\n",
              "      <td>Seller</td>\n",
              "      <td>PARTY</td>\n",
              "      <td>575</td>\n",
              "      <td>588</td>\n",
              "      <td>Buyer Entities</td>\n",
              "      <td>0.5416529</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>73</th>\n",
              "      <td>has_collective_alias</td>\n",
              "      <td>PARTY</td>\n",
              "      <td>293</td>\n",
              "      <td>298</td>\n",
              "      <td>Seller</td>\n",
              "      <td>ALIAS</td>\n",
              "      <td>612</td>\n",
              "      <td>616</td>\n",
              "      <td>Party</td>\n",
              "      <td>0.5416529</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>74</th>\n",
              "      <td>has_collective_alias</td>\n",
              "      <td>PARTY</td>\n",
              "      <td>293</td>\n",
              "      <td>298</td>\n",
              "      <td>Seller</td>\n",
              "      <td>ALIAS</td>\n",
              "      <td>642</td>\n",
              "      <td>648</td>\n",
              "      <td>Parties</td>\n",
              "      <td>0.5416529</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>82</th>\n",
              "      <td>has_collective_alias</td>\n",
              "      <td>ALIAS</td>\n",
              "      <td>400</td>\n",
              "      <td>404</td>\n",
              "      <td>Buyer</td>\n",
              "      <td>ALIAS</td>\n",
              "      <td>479</td>\n",
              "      <td>485</td>\n",
              "      <td>Company</td>\n",
              "      <td>0.60943455</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>83</th>\n",
              "      <td>has_collective_alias</td>\n",
              "      <td>ALIAS</td>\n",
              "      <td>400</td>\n",
              "      <td>404</td>\n",
              "      <td>Buyer</td>\n",
              "      <td>ALIAS</td>\n",
              "      <td>517</td>\n",
              "      <td>530</td>\n",
              "      <td>Buyer Entities</td>\n",
              "      <td>0.60943455</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>84</th>\n",
              "      <td>has_collective_alias</td>\n",
              "      <td>ALIAS</td>\n",
              "      <td>400</td>\n",
              "      <td>404</td>\n",
              "      <td>Buyer</td>\n",
              "      <td>PARTY</td>\n",
              "      <td>575</td>\n",
              "      <td>588</td>\n",
              "      <td>Buyer Entities</td>\n",
              "      <td>0.60943455</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>85</th>\n",
              "      <td>has_collective_alias</td>\n",
              "      <td>ALIAS</td>\n",
              "      <td>400</td>\n",
              "      <td>404</td>\n",
              "      <td>Buyer</td>\n",
              "      <td>ALIAS</td>\n",
              "      <td>612</td>\n",
              "      <td>616</td>\n",
              "      <td>Party</td>\n",
              "      <td>0.60943455</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>86</th>\n",
              "      <td>has_collective_alias</td>\n",
              "      <td>ALIAS</td>\n",
              "      <td>400</td>\n",
              "      <td>404</td>\n",
              "      <td>Buyer</td>\n",
              "      <td>ALIAS</td>\n",
              "      <td>642</td>\n",
              "      <td>648</td>\n",
              "      <td>Parties</td>\n",
              "      <td>0.60943455</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>\n",
              "      <button class=\"colab-df-convert\" onclick=\"convertToInteractive('df-f6d50ef3-a2ba-4632-91dc-b5dd7423352d')\"\n",
              "              title=\"Convert this dataframe to an interactive table.\"\n",
              "              style=\"display:none;\">\n",
              "        \n",
              "  <svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\"viewBox=\"0 0 24 24\"\n",
              "       width=\"24px\">\n",
              "    <path d=\"M0 0h24v24H0V0z\" fill=\"none\"/>\n",
              "    <path d=\"M18.56 5.44l.94 2.06.94-2.06 2.06-.94-2.06-.94-.94-2.06-.94 2.06-2.06.94zm-11 1L8.5 8.5l.94-2.06 2.06-.94-2.06-.94L8.5 2.5l-.94 2.06-2.06.94zm10 10l.94 2.06.94-2.06 2.06-.94-2.06-.94-.94-2.06-.94 2.06-2.06.94z\"/><path d=\"M17.41 7.96l-1.37-1.37c-.4-.4-.92-.59-1.43-.59-.52 0-1.04.2-1.43.59L10.3 9.45l-7.72 7.72c-.78.78-.78 2.05 0 2.83L4 21.41c.39.39.9.59 1.41.59.51 0 1.02-.2 1.41-.59l7.78-7.78 2.81-2.81c.8-.78.8-2.07 0-2.86zM5.41 20L4 18.59l7.72-7.72 1.47 1.35L5.41 20z\"/>\n",
              "  </svg>\n",
              "      </button>\n",
              "      \n",
              "  <style>\n",
              "    .colab-df-container {\n",
              "      display:flex;\n",
              "      flex-wrap:wrap;\n",
              "      gap: 12px;\n",
              "    }\n",
              "\n",
              "    .colab-df-convert {\n",
              "      background-color: #E8F0FE;\n",
              "      border: none;\n",
              "      border-radius: 50%;\n",
              "      cursor: pointer;\n",
              "      display: none;\n",
              "      fill: #1967D2;\n",
              "      height: 32px;\n",
              "      padding: 0 0 0 0;\n",
              "      width: 32px;\n",
              "    }\n",
              "\n",
              "    .colab-df-convert:hover {\n",
              "      background-color: #E2EBFA;\n",
              "      box-shadow: 0px 1px 2px rgba(60, 64, 67, 0.3), 0px 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
              "      fill: #174EA6;\n",
              "    }\n",
              "\n",
              "    [theme=dark] .colab-df-convert {\n",
              "      background-color: #3B4455;\n",
              "      fill: #D2E3FC;\n",
              "    }\n",
              "\n",
              "    [theme=dark] .colab-df-convert:hover {\n",
              "      background-color: #434B5C;\n",
              "      box-shadow: 0px 1px 3px 1px rgba(0, 0, 0, 0.15);\n",
              "      filter: drop-shadow(0px 1px 2px rgba(0, 0, 0, 0.3));\n",
              "      fill: #FFFFFF;\n",
              "    }\n",
              "  </style>\n",
              "\n",
              "      <script>\n",
              "        const buttonEl =\n",
              "          document.querySelector('#df-f6d50ef3-a2ba-4632-91dc-b5dd7423352d button.colab-df-convert');\n",
              "        buttonEl.style.display =\n",
              "          google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
              "\n",
              "        async function convertToInteractive(key) {\n",
              "          const element = document.querySelector('#df-f6d50ef3-a2ba-4632-91dc-b5dd7423352d');\n",
              "          const dataTable =\n",
              "            await google.colab.kernel.invokeFunction('convertToInteractive',\n",
              "                                                     [key], {});\n",
              "          if (!dataTable) return;\n",
              "\n",
              "          const docLinkHtml = 'Like what you see? Visit the ' +\n",
              "            '<a target=\"_blank\" href=https://colab.research.google.com/notebooks/data_table.ipynb>data table notebook</a>'\n",
              "            + ' to learn more about interactive tables.';\n",
              "          element.innerHTML = '';\n",
              "          dataTable['output_type'] = 'display_data';\n",
              "          await google.colab.output.renderOutput(dataTable, element);\n",
              "          const docLink = document.createElement('div');\n",
              "          docLink.innerHTML = docLinkHtml;\n",
              "          element.appendChild(docLink);\n",
              "        }\n",
              "      </script>\n",
              "    </div>\n",
              "  </div>\n",
              "  "
            ]
          },
          "metadata": {},
          "execution_count": 98
        }
      ],
      "source": [
        "rel_df = get_relations_df(results)\n",
        "rel_df = rel_df[rel_df['relation']!='other']\n",
        "rel_df"
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "i50VHOpXRIa0"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "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.7.15"
    },
    "varInspector": {
      "cols": {
        "lenName": 16,
        "lenType": 16,
        "lenVar": 40
      },
      "kernels_config": {
        "python": {
          "delete_cmd_postfix": "",
          "delete_cmd_prefix": "del ",
          "library": "var_list.py",
          "varRefreshCmd": "print(var_dic_list())"
        },
        "r": {
          "delete_cmd_postfix": ") ",
          "delete_cmd_prefix": "rm(",
          "library": "var_list.r",
          "varRefreshCmd": "cat(var_dic_list()) "
        }
      },
      "types_to_exclude": [
        "module",
        "function",
        "builtin_function_or_method",
        "instance",
        "_Feature"
      ],
      "window_display": false
    },
    "colab": {
      "provenance": []
    },
    "gpuClass": "standard",
    "accelerator": "GPU"
  },
  "nbformat": 4,
  "nbformat_minor": 0
}