title named entity recognition using transformers varun singhhttps www linkedin cominvarunsingh2 date created 2021623 last modified 2023725 description ner using the transformers and data from conll 2003 shared task accelerator gpu introduction named entity recognition ner is the process of identifying named entities in text example of named entities are person location organization dates etc ner is essentially a token classification task where every token is classified into one or more predetermined categories in this exercise we will train a simple transformer based model to perform ner we will be using the data from conll 2003 shared task for more information about the dataset please visit the dataset websitehttps www clips uantwerpen beconll2003ner however since obtaining this data requires an additional step of getting a free license we will be using huggingface s datasets library which contains a processed version of this dataset install the open source datasets library from huggingface we also download the script used to evaluate ner models shell pip3 install datasets wget https raw githubusercontent comsighsmileconllevalmasterconlleval py imports for data preprocessing next let s define a tokenandpositionembedding layer build the ner model class as a keras model subclass load the conll 2003 dataset from the datasets library and process it we will export this data to a tabseparated file format which will be easy to read as a tf data dataset object make the ner label lookup table ner labels are usually provided in iob iob2 or iobes formats checkout this link for more information wikipediahttps en wikipedia orgwikiinsidee28093outsidee28093beginningtagging note that we start our label numbering from 1 since 0 will be reserved for padding we have a total of 10 labels 9 from the ner dataset and one for padding get a list of all tokens in the training dataset this will be used to create the vocabulary we only take vocabsize 2 most commons words from the training data since the stringlookup class uses 2 additional tokens one denoting an unknown token and another one denoting a masking token the stringlook class will convert tokens to token ids create 2 new dataset objects from the training and validation data print out one line to make sure it looks good the first record in the line is the number of tokens after that we will have all the tokens followed by all the ner tags we will be using the following map function to transform the data in the dataset data preprocessing using tensorflow we use paddedbatch here because each record in the dataset has a different length we will be using a custom loss function that will ignore the loss from padded tokens compile and fit the model sample inference using the trained model eu borg german bmisc british bmisc metrics calculation here is a function to calculate the metrics the function calculates f1 score for the overall ner dataset as well as individual scores for each ner tag conclusions in this exercise we created a simple transformer based named entity recognition model we trained it on the conll 2003 shared task data and got an overall f1 score of around 70 state of the art ner models finetuned on pretrained models such as bert or electra can easily get much higher f1 score between 9095 on this dataset owing to the inherent knowledge of words as part of the pretraining process and the usage of subword tokenization introduction named entity recognition ner is the process of identifying named entities in text example of named entities are person location organization dates etc ner is essentially a token classification task where every token is classified into one or more predetermined categories in this exercise we will train a simple transformer based model to perform ner we will be using the data from conll 2003 shared task for more information about the dataset please visit the dataset website https www clips uantwerpen be conll2003 ner however since obtaining this data requires an additional step of getting a free license we will be using huggingface s datasets library which contains a processed version of this dataset install the open source datasets library from huggingface we also download the script used to evaluate ner models shell pip3 install datasets wget https raw githubusercontent com sighsmile conlleval master conlleval py imports for data preprocessing next let s define a tokenandpositionembedding layer build the ner model class as a keras model subclass load the conll 2003 dataset from the datasets library and process it we will export this data to a tab separated file format which will be easy to read as a tf data dataset object make the ner label lookup table ner labels are usually provided in iob iob2 or iobes formats checkout this link for more information wikipedia https en wikipedia org wiki inside e2 80 93outside e2 80 93beginning_ tagging note that we start our label numbering from 1 since 0 will be reserved for padding we have a total of 10 labels 9 from the ner dataset and one for padding get a list of all tokens in the training dataset this will be used to create the vocabulary we only take vocab_size 2 most commons words from the training data since the stringlookup class uses 2 additional tokens one denoting an unknown token and another one denoting a masking token the stringlook class will convert tokens to token ids create 2 new dataset objects from the training and validation data print out one line to make sure it looks good the first record in the line is the number of tokens after that we will have all the tokens followed by all the ner tags we will be using the following map function to transform the data in the dataset data preprocessing using tensorflow we use padded_batch here because each record in the dataset has a different length we will be using a custom loss function that will ignore the loss from padded tokens compile and fit the model sample inference using the trained model eu b org german b misc british b misc metrics calculation here is a function to calculate the metrics the function calculates f1 score for the overall ner dataset as well as individual scores for each ner tag conclusions in this exercise we created a simple transformer based named entity recognition model we trained it on the conll 2003 shared task data and got an overall f1 score of around 70 state of the art ner models fine tuned on pretrained models such as bert or electra can easily get much higher f1 score between 90 95 on this dataset owing to the inherent knowledge of words as part of the pretraining process and the usage of subword tokenization
import os
import numpy as np
import keras
from keras import layers
from datasets import load_dataset
from collections import Counter
from conlleval import evaluate


from tensorflow import data as tf_data
from tensorflow import strings as tf_strings


class TransformerBlock(layers.Layer):
    def __init__(self, embed_dim, num_heads, ff_dim, rate=0.1):
        super().__init__()
        self.att = keras.layers.MultiHeadAttention(
            num_heads=num_heads, key_dim=embed_dim
        )
        self.ffn = keras.Sequential(
            [
                keras.layers.Dense(ff_dim, activation="relu"),
                keras.layers.Dense(embed_dim),
            ]
        )
        self.layernorm1 = keras.layers.LayerNormalization(epsilon=1e-6)
        self.layernorm2 = keras.layers.LayerNormalization(epsilon=1e-6)
        self.dropout1 = keras.layers.Dropout(rate)
        self.dropout2 = keras.layers.Dropout(rate)

    def call(self, inputs, training=False):
        attn_output = self.att(inputs, inputs)
        attn_output = self.dropout1(attn_output, training=training)
        out1 = self.layernorm1(inputs + attn_output)
        ffn_output = self.ffn(out1)
        ffn_output = self.dropout2(ffn_output, training=training)
        return self.layernorm2(out1 + ffn_output)





class TokenAndPositionEmbedding(layers.Layer):
    def __init__(self, maxlen, vocab_size, embed_dim):
        super().__init__()
        self.token_emb = keras.layers.Embedding(
            input_dim=vocab_size, output_dim=embed_dim
        )
        self.pos_emb = keras.layers.Embedding(
            input_dim=maxlen, output_dim=embed_dim
        )

    def call(self, inputs):
        maxlen = keras.ops.backend.shape(inputs)[-1]
        positions = keras.ops.arange(start=0, stop=maxlen, step=1)
        position_embeddings = self.pos_emb(positions)
        token_embeddings = self.token_emb(inputs)
        return token_embeddings + position_embeddings





class NERModel(keras.Model):
    def __init__(
        self,
        num_tags,
        vocab_size,
        maxlen=128,
        embed_dim=32,
        num_heads=2,
        ff_dim=32,
    ):
        super().__init__()
        self.embedding_layer = TokenAndPositionEmbedding(
            maxlen, vocab_size, embed_dim
        )
        self.transformer_block = TransformerBlock(embed_dim, num_heads, ff_dim)
        self.dropout1 = layers.Dropout(0.1)
        self.ff = layers.Dense(ff_dim, activation="relu")
        self.dropout2 = layers.Dropout(0.1)
        self.ff_final = layers.Dense(num_tags, activation="softmax")

    def call(self, inputs, training=False):
        x = self.embedding_layer(inputs)
        x = self.transformer_block(x)
        x = self.dropout1(x, training=training)
        x = self.ff(x)
        x = self.dropout2(x, training=training)
        x = self.ff_final(x)
        return x




conll_data = load_dataset("conll2003")




def export_to_file(export_file_path, data):
    with open(export_file_path, "w") as f:
        for record in data:
            ner_tags = record["ner_tags"]
            tokens = record["tokens"]
            if len(tokens) > 0:
                f.write(
                    str(len(tokens))
                    + "\t"
                    + "\t".join(tokens)
                    + "\t"
                    + "\t".join(map(str, ner_tags))
                    + "\n"
                )


os.mkdir("data")
export_to_file("./data/conll_train.txt", conll_data["train"])
export_to_file("./data/conll_val.txt", conll_data["validation"])




def make_tag_lookup_table():
    iob_labels = ["B", "I"]
    ner_labels = ["PER", "ORG", "LOC", "MISC"]
    all_labels = [
        (label1, label2) for label2 in ner_labels for label1 in iob_labels
    ]
    all_labels = ["-".join([a, b]) for a, b in all_labels]
    all_labels = ["[PAD]", "O"] + all_labels
    return dict(zip(range(0, len(all_labels) + 1), all_labels))


mapping = make_tag_lookup_table()
print(mapping)



all_tokens = sum(conll_data["train"]["tokens"], [])
all_tokens_array = np.array(list(map(str.lower, all_tokens)))

counter = Counter(all_tokens_array)
print(len(counter))

num_tags = len(mapping)
vocab_size = 20000




vocabulary = [token for token, count in counter.most_common(vocab_size - 2)]


lookup_layer = keras.layers.StringLookup(vocabulary=vocabulary)



train_data = tf_data.TextLineDataset("./data/conll_train.txt")
val_data = tf_data.TextLineDataset("./data/conll_val.txt")



print(list(train_data.take(1).as_numpy_iterator()))





def map_record_to_training_data(record):
    record = tf_strings.split(record, sep="\t")
    length = tf_strings.to_number(record[0], out_type="int32")
    tokens = record[1 : length + 1]
    tags = record[length + 1 :]
    tags = tf_strings.to_number(tags, out_type="int64")
    tags += 1
    return tokens, tags


def lowercase_and_convert_to_ids(tokens):
    tokens = tf_strings.lower(tokens)
    return lookup_layer(tokens)




batch_size = 32
train_dataset = (
    train_data.map(map_record_to_training_data)
    .map(lambda x, y: (lowercase_and_convert_to_ids(x), y))
    .padded_batch(batch_size)
)
val_dataset = (
    val_data.map(map_record_to_training_data)
    .map(lambda x, y: (lowercase_and_convert_to_ids(x), y))
    .padded_batch(batch_size)
)

ner_model = NERModel(num_tags, vocab_size, embed_dim=32, num_heads=4, ff_dim=64)




class CustomNonPaddingTokenLoss(keras.losses.Loss):
    def __init__(self, name="custom_ner_loss"):
        super().__init__(name=name)

    def call(self, y_true, y_pred):
        loss_fn = keras.losses.SparseCategoricalCrossentropy(
            from_logits=True, reduction=None
        )
        loss = loss_fn(y_true, y_pred)
        mask = keras.backend.cast((y_true > 0), dtype="float32")
        loss = loss * mask
        return keras.ops.sum(loss) / keras.ops.sum(mask)


loss = CustomNonPaddingTokenLoss()



ner_model.compile(optimizer="adam", loss=loss)
ner_model.fit(train_dataset, epochs=10)


def tokenize_and_convert_to_ids(text):
    tokens = text.split()
    return lowercase_and_convert_to_ids(tokens)



sample_input = tokenize_and_convert_to_ids(
    "eu rejects german call to boycott british lamb"
)
sample_input = keras.ops.reshape(sample_input, new_shape=[1, -1])
print(sample_input)

output = ner_model.predict(sample_input)
prediction = np.argmax(output, axis=-1)[0]
prediction = [mapping[i] for i in prediction]


print(prediction)




def calculate_metrics(dataset):
    all_true_tag_ids, all_predicted_tag_ids = [], []

    for x, y in dataset:
        output = ner_model.predict(x)
        predictions = np.argmax(output, axis=-1)
        predictions = np.reshape(predictions, [-1])

        true_tag_ids = np.reshape(y, [-1])

        mask = (true_tag_ids > 0) & (predictions > 0)
        true_tag_ids = true_tag_ids[mask]
        predicted_tag_ids = predictions[mask]

        all_true_tag_ids.append(true_tag_ids)
        all_predicted_tag_ids.append(predicted_tag_ids)

    all_true_tag_ids = np.concatenate(all_true_tag_ids)
    all_predicted_tag_ids = np.concatenate(all_predicted_tag_ids)

    predicted_tags = [mapping[tag] for tag in all_predicted_tag_ids]
    real_tags = [mapping[tag] for tag in all_true_tag_ids]

    evaluate(real_tags, predicted_tags)


calculate_metrics(val_dataset)