Document Understanding Solution - Relationship Extraction

This notebook's CI test result for us-west-2 is as follows. CI test results in other regions can be found at the end of the notebook.

Relation Extraction (RE) is the task of extracting semantic relationships from text, which usually occur between two or more entities. In this notebook, we demonstrate two use cases of Relation Extraction:

1. How to fine-tune a pre-trained Transformer model on a custom dataset, and then run inference on the fine-tuned model.
2. How to run SageMaker Automatic Model Tuning (a hyperparameter optimization procedure) to find the best model compared with the model fine-tuned in point 1. The performance of the optimal model and model fine-tuned in point 1 is evaluated on a hold-out test data.

Note: When running this notebook on SageMaker Studio, you should make sure the PyTorch 1.10 Python 3.8 CPU Optimized image/kernel is used. When running this notebook on SageMaker Notebook Instance, you should make sure the 'sagemaker-soln' kernel is used.

This solution relies on a config file to run the provisioned AWS resources. Run the cell below to generate that file.

1. Set Up

We start by importing a variety of packages that are used throughout the notebook. One of the most important packages is the Amazon SageMaker Python SDK (i.e. import sagemaker). We also import modules from our own custom (and editable) package that can be found at ../package.

2. Finetune the pre-trained model on a custom dataset

This is a Relationship Extraction model built on a Bert-base-uncased using transformers from the transformers library.

The model for fine-tuning attaches a linear classification layer that takes a pair of token embeddings outputted by the Text Embedding model and initializes the layer parameters to random values. The fine-tuning step fine-tunes all the model parameters to minimize prediction error on the input data and returns the fine-tuned model. The Text Embedding model we use in this demonstartion is Bert-base-uncased from the transformers library. The dataset we fine-tune the model is SemEval-2010 Task 8. The SemEval-2 Task 8 is a dataset for multi-way classification of mutually exclusive semantic relations between pairs of nominals.

The model returned by fine-tuning can be further deployed for inference. Below are the instructions for how the training data should be formatted for input to the model.

• Input: A directory containing a txt format file.
  • Each observation contains three components, text, semantic relation label, and comment (optional), each of which takes a line in the txt format file. Observations are separated by an empty line. For each observation, there are markers highlighting the two terms in the text and their semantic relation label in the line below.
• Output: A trained model that can be deployed for inference.

Below is an example of txt format file. Note. Desipte of the same semantic relation label, pairs of entities with different order relations are counted as different labels. For an example, Component-Whole(e2,e1) and Component-Whole(e1,e2) are different semantic relation labels. The data for training and validation are downloaded into directory ../data/semeval2010t8 in the following section.

Citation: @inproceedings{hendrickx-etal-2010-semeval, title = "{S}em{E}val-2010 Task 8: Multi-Way Classification of Semantic Relations between Pairs of Nominals", author = "Hendrickx, Iris and Kim, Su Nam and Kozareva, Zornitsa and Nakov, Preslav and {'O} S{'e}aghdha, Diarmuid and Pad{'o}, Sebastian and Pennacchiotti, Marco and Romano, Lorenza and Szpakowicz, Stan", booktitle = "Proceedings of the 5th International Workshop on Semantic Evaluation", month = jul, year = "2010", address = "Uppsala, Sweden", publisher = "Association for Computational Linguistics", url = "https://www.aclweb.org/anthology/S10-1006", pages = "33--38", }

2.1. Download, preprocess, and upload the training data

The dataset has been partitioned into train.txt, validation.txt, and test.txt data. Thus we don't need split the train data as what we do in previous notebooks. Thetrain.txt and validation.txt are used as training and validation data. The test.txt is used as hold-out test data to evaluate model performance with / without hyperparameter optimization. Next, we upload them into S3 path which are used as input for training.

2.2. Set Training parameters

Now that we are done with all the setup that is needed, we are ready to fine-tune our relation extraction model.

3.2. Fine-tuning without hyperparameter optimization

We use the PyTorch from the Amazon SageMaker Python SDK. The entry script is located under ../containers/relationship_extraction/entry_point.py

3.3. Deploy & run Inference on the fine-tuned model

A trained model does nothing on its own. We now want to use the model to perform inference. For this example, it means predicting the semantic relation label of two text string within an input text.

We use the unique solution prefix to name the model and endpoint.

When calling our new endpoint from the notebook, we use a Amazon SageMaker SDK Predictor. A Predictor is used to send data to an endpoint (as part of a request), and interpret the response. Our estimator.deploy command returned a Predictor but, by default, it sends and receive numpy arrays. Our endpoint expects to receive (and also sends) JSON formatted objects, so we modify the Predictor to use JSON instead of the PyTorch endpoint default of numpy arrays. JSON is used here because it is a standard endpoint format and the endpoint response can contain nested data structures.

With our model successfully deployed and our predictor configured, we can try out the relationship extraction model out on example inputs.

Next, let's query the deployed endpoint to get for the prediction for each test example located in ../data/semeval2010t8/test/test.txt.

Since the task is essentially multiclass classification task, we use accuracy, f1 macro, and f1 micro as the evaluation scores. For each of them, higher value indicates better results.

3. Finetune the pre-trained model on a custom dataset with automatic model tuning (AMT)

Amazon SageMaker automatic model tuning, also known as hyperparameter tuning, finds the best version of a model by running many training jobs on your dataset using the algorithm and ranges of hyperparameters that you specify. It then chooses the hyperparameter values that result in a model that performs the best, as measured by a metric that you choose. We use a HyperparameterTuner object to interact with Amazon SageMaker hyperparameter tuning APIs.

3.1. Fine-tuning with hyperparameter optimization

3.2. Deploy & run Inference on the fine-tuned model

We can see results with hyperparameter optimization shows better performance on the hold-out test data.

3.3. Clean Up the endpoint

When you've finished with the summarization endpoint (and associated endpoint-config), make sure that you delete it to avoid accidental charges.

Notebook CI Test Results

This notebook was tested in multiple regions. The test results are as follows, except for us-west-2 which is shown at the top of the notebook.