pos_tagger.py

"""
POS tagger using fine-tuned BERT
Most of the code in this module is borrowed from https://github.com/soutsios/pos-tagger-bert
"""
import os
from pathlib import Path

FILE_DIR = os.path.dirname(os.path.realpath(__file__))
STATIC_DIR = os.path.join(FILE_DIR, "static/converter")
STATIC_DIR = Path(STATIC_DIR)
# trained model weights for POS
MODEL_WEIGHTS_DIR = Path(f"{STATIC_DIR}/bert_model/bert_last_epoch.h5")


# Check if the file with the weights exists
if not os.path.exists(MODEL_WEIGHTS_DIR):
    # if not, only expose a tag_pos function that doesn't do anything
    print(
        f"No model weights found at {MODEL_WEIGHTS_DIR}. Converter will not do POS tagging."
    )

    def tag_pos(t):
        return None


else:
    # if the weigts are there, initialise the entire BERT model
    print(
        f"Model weights found at {MODEL_WEIGHTS_DIR}. Loading model for POS tagging..."
    )

    import keras
    import numpy as np

    from keras.layers import Layer
    from keras import backend as K

    import tensorflow as tf
    import tensorflow_hub as hub
    from bert.tokenization import FullTokenizer

    from functools import lru_cache

    # from tqdm import tqdm_notebook

    class PaddingInputExample(object):
        """Fake example so the num input examples is a multiple of the batch size.
        When running eval/predict on the TPU, we need to pad the number of examples
        to be a multiple of the batch size, because the TPU requires a fixed batch
        size. The alternative is to drop the last batch, which is bad because it means
        the entire output data won't be generated.
        We use this class instead of `None` because treating `None` as padding
        battches could cause silent errors.
        """

    class InputExample(object):
        """A single training/test example for simple sequence classification."""

        def __init__(self, guid, text_a, text_b=None, label=None):
            """Constructs a InputExample.
            Args:
            guid: Unique id for the example.
            text_a: string. The untokenized text of the first sequence. For single
                sequence tasks, only this sequence must be specified.
            text_b: (Optional) string. The untokenized text of the second sequence.
                Only must be specified for sequence pair tasks.
            label: (Optional) string. The label of the example. This should be
                specified for train and dev examples, but not for test examples.
            """
            self.guid = guid
            self.text_a = text_a
            self.text_b = text_b
            self.label = label

    def create_tokenizer_from_hub_module(bert_path):
        """Get the vocab file and casing info from the Hub module."""
        bert_module = hub.Module(bert_path)
        tokenization_info = bert_module(signature="tokenization_info", as_dict=True)
        vocab_file, do_lower_case = sess.run(
            [
                tokenization_info["vocab_file"],
                tokenization_info["do_lower_case"],
            ]
        )
        return FullTokenizer(
            vocab_file=vocab_file, do_lower_case=do_lower_case
        )  # , spm_model_file=vocab_file

    def convert_single_example(tokenizer, example, max_seq_length=256):
        """Converts a single `InputExample` into a single `InputFeatures`."""

        if isinstance(example, PaddingInputExample):
            input_ids = [0] * max_seq_length
            input_mask = [0] * max_seq_length
            segment_ids = [0] * max_seq_length
            label_ids = [0] * max_seq_length
            return input_ids, input_mask, segment_ids, label_ids

        tokens_a = example.text_a
        if len(tokens_a) > max_seq_length - 2:
            tokens_a = tokens_a[0 : (max_seq_length - 2)]

        # Token map will be an int -> int mapping between the `orig_tokens` index and
        # the `bert_tokens` index.

        # bert_tokens == ["[CLS]", "john", "johan", "##son", "'", "s", "house", "[SEP]"]
        # orig_to_tok_map == [1, 2, 4, 6]
        orig_to_tok_map = []
        tokens = []
        segment_ids = []

        tokens.append("[CLS]")
        segment_ids.append(0)
        orig_to_tok_map.append(len(tokens) - 1)
        # print(len(tokens_a))
        for token in tokens_a:
            tokens.extend(tokenizer.tokenize(token))
            orig_to_tok_map.append(len(tokens) - 1)
            segment_ids.append(0)
        tokens.append("[SEP]")
        segment_ids.append(0)
        orig_to_tok_map.append(len(tokens) - 1)
        input_ids = tokenizer.convert_tokens_to_ids(
            [tokens[i] for i in orig_to_tok_map]
        )
        # print(len(orig_to_tok_map), len(tokens), len(input_ids), len(segment_ids)) #for debugging

        # The mask has 1 for real tokens and 0 for padding tokens. Only real
        # tokens are attended to.
        input_mask = [1] * len(input_ids)

        label_ids = []
        labels = example.label
        label_ids.append(0)
        label_ids.extend([tag2int[label] for label in labels])
        label_ids.append(0)
        # print(len(label_ids)) #for debugging
        # Zero-pad up to the sequence length.
        while len(input_ids) < max_seq_length:
            input_ids.append(0)
            input_mask.append(0)
            segment_ids.append(0)
            label_ids.append(0)

        assert len(input_ids) == max_seq_length
        assert len(input_mask) == max_seq_length
        assert len(segment_ids) == max_seq_length
        assert len(label_ids) == max_seq_length

        return input_ids, input_mask, segment_ids, label_ids

    def convert_examples_to_features(tokenizer, examples, max_seq_length=256):
        """Convert a set of `InputExample`s to a list of `InputFeatures`."""

        input_ids, input_masks, segment_ids, labels = [], [], [], []
        # tqdm_notebook(examples, desc="Converting examples to features"):
        for example in examples:
            input_id, input_mask, segment_id, label = convert_single_example(
                tokenizer, example, max_seq_length
            )
            input_ids.append(input_id)
            input_masks.append(input_mask)
            segment_ids.append(segment_id)
            labels.append(label)
        return (
            np.array(input_ids),
            np.array(input_masks),
            np.array(segment_ids),
            np.array(labels),
        )

    def convert_text_to_examples(texts, labels):
        """Create InputExamples"""
        InputExamples = []
        for text, label in zip(texts, labels):
            InputExamples.append(
                InputExample(guid=None, text_a=text, text_b=None, label=label)
            )
        return InputExamples

    class BertLayer(Layer):
        def __init__(
            self, output_representation="sequence_output", trainable=True, **kwargs
        ):
            self.bert = None
            super(BertLayer, self).__init__(**kwargs)

            self.trainable = trainable
            self.output_representation = output_representation

        def build(self, input_shape):
            # SetUp tensorflow Hub module
            self.bert = hub.Module(
                bert_path, trainable=self.trainable, name="{}_module".format(self.name)
            )

            # Assign module's trainable weights to model
            # Remove unused layers and set trainable parameters
            # s = ["/cls/", "/pooler/", 'layer_11', 'layer_10', 'layer_9', 'layer_8', 'layer_7', 'layer_6']
            s = ["/cls/", "/pooler/"]
            self.trainable_weights += [
                var
                for var in self.bert.variables[:]
                if not any(x in var.name for x in s)
            ]

            for var in self.bert.variables:
                if var not in self._trainable_weights:
                    self._non_trainable_weights.append(var)

            # See Trainable Variables
            # tf.logging.info("**** Trainable Variables ****")
            # for var in self.trainable_weights:
            #    init_string = ", *INIT_FROM_CKPT*"
            #    tf.logging.info("  name = %s, shape = %s%s", var.name, var.shape, init_string)

            print("Trainable weights:", len(self.trainable_weights))
            super(BertLayer, self).build(input_shape)

        def call(self, inputs, mask=None):
            inputs = [K.cast(x, dtype="int32") for x in inputs]
            input_ids, input_mask, segment_ids = inputs
            bert_inputs = dict(
                input_ids=input_ids, input_mask=input_mask, segment_ids=segment_ids
            )
            result = self.bert(inputs=bert_inputs, signature="tokens", as_dict=True)[
                self.output_representation
            ]
            return result

        def compute_mask(self, inputs, mask=None):
            return K.not_equal(inputs[0], 0.0)

        def compute_output_shape(self, input_shape):
            if self.output_representation == "pooled_output":
                return (None, 768)
            else:
                return (None, None, 768)

    # Build model
    def build_model(max_seq_length):
        seed = 0
        in_id = keras.layers.Input(shape=(max_seq_length,), name="input_ids")
        in_mask = keras.layers.Input(shape=(max_seq_length,), name="input_masks")
        in_segment = keras.layers.Input(shape=(max_seq_length,), name="segment_ids")
        bert_inputs = [in_id, in_mask, in_segment]

        np.random.seed(seed)
        bert_output = BertLayer()(bert_inputs)

        np.random.seed(seed)
        outputs = keras.layers.Dense(n_tags, activation=keras.activations.softmax)(
            bert_output
        )

        np.random.seed(seed)
        model = keras.models.Model(inputs=bert_inputs, outputs=outputs)
        np.random.seed(seed)
        model.compile(
            optimizer=keras.optimizers.Adam(lr=0.000008),
            loss=keras.losses.categorical_crossentropy,
            metrics=["accuracy"],
        )
        model.summary(100)
        return model

    def initialize_vars(sess):
        sess.run(tf.compat.v1.local_variables_initializer())
        sess.run(tf.compat.v1.global_variables_initializer())
        sess.run(tf.compat.v1.tables_initializer())
        K.set_session(sess)

    def split(sentences, max):
        """
        Split the sentences to MAX_SEQUENCE_LENGTH and so the number of samples increases accordingly.
        For example, if MAX_SEQUENCE_LENGTH=70, a sentence with length 150 splits in 3 sentences: 150=70+70+10
        """
        new = []
        for data in sentences:
            new.append(([data[x : x + max] for x in range(0, len(data), max)]))
        new = [val for sublist in new for val in sublist]
        return new

    @lru_cache(maxsize=42)
    def tag_pos(sentence_string):
        """Do POS tagging for each word in the provided sentence
        (This function is cached because it's very CPU intensive and it might be often that we want to call it for words in the same sentence)

        Args:
            sentence_string (str): Sentence on the words of which we want to do POS tagging

        Returns:
            list: List of words and their respective POS tags
        """

        sentence_tokenized = sentence_string.split(" ")

        # split into multiple sentences of max length, if words in input exeed max lenght
        if len(sentence_tokenized) > MAX_SEQUENCE_LENGTH:
            sentence_tokenized = split([sentence_tokenized], MAX_SEQUENCE_LENGTH)
        else:
            sentence_tokenized = [sentence_tokenized]

        # where the results will be stored
        pred_tuples = []

        for sentence_ini in sentence_tokenized:

            tokens_a = sentence_ini

            orig_to_tok_map = []
            tokens = []
            segment_ids = []
            tokens.append("[CLS]")
            segment_ids.append(0)
            orig_to_tok_map.append(len(tokens) - 1)
            for token in tokens_a:
                # orig_to_tok_map.append(len(tokens)) # keep first piece of tokenized term
                tokens.extend(tokenizer.tokenize(token))
                orig_to_tok_map.append(
                    len(tokens) - 1
                )  # # keep last piece of tokenized term -->> gives better results!
                segment_ids.append(0)
            tokens.append("[SEP]")
            segment_ids.append(0)
            orig_to_tok_map.append(len(tokens) - 1)
            input_ids = tokenizer.convert_tokens_to_ids(
                [tokens[i] for i in orig_to_tok_map]
            )

            # Convert data to InputExample format
            test_example = convert_text_to_examples(
                [sentence_ini], [["-PAD-"] * len(sentence_ini)]
            )

            # Convert to features
            (input_ids, input_masks, segment_ids, _) = convert_examples_to_features(
                tokenizer, test_example, max_seq_length=MAX_SEQUENCE_LENGTH + 2
            )

            with sess.as_default():
                with sess.graph.as_default():
                    predictions = model.predict(
                        [input_ids, input_masks, segment_ids], batch_size=1
                    ).argmax(-1)[0]

            pred_tuples += [
                (sentence_ini[i - 1], int2tag[pred])
                for i, pred in enumerate(predictions)
                if not i > len(sentence_ini) and pred != 0
            ]

        return pred_tuples

    MAX_SEQUENCE_LENGTH = 70

    # Initialize session
    tf_config = tf.compat.v1.ConfigProto(allow_soft_placement=True)
    tf_config.gpu_options.allow_growth = True
    # tf_config.intra_op_parallelism_threads = 2
    # tf_config.inter_op_parallelism_threads = 2
    sess = tf.compat.v1.Session(config=tf_config)
    initialize_vars(sess)

    tags = {
        "ADJ",
        "ADP",
        "ADV",
        "AUX",
        "CCONJ",
        "DET",
        "INTJ",
        "NOUN",
        "NUM",
        "PART",
        "PRON",
        "PROPN",
        "PUNCT",
        "SCONJ",
        "VERB",
        "X",
    }

    # label mappings
    tag2int = {}
    int2tag = {}

    for i, tag in enumerate(sorted(tags)):
        tag2int[tag] = i + 1
        int2tag[i + 1] = tag

    # Special character for the tags
    tag2int["-PAD-"] = 0
    int2tag[0] = "-PAD-"

    n_tags = len(tag2int)

    # graph = tf.get_default_graph()

    # Params for bert model and tokenization
    bert_path = "https://tfhub.dev/google/bert_multi_cased_L-12_H-768_A-12/1"  # use multi lang version!

    # Instantiate tokenizer
    tokenizer = create_tokenizer_from_hub_module(bert_path)

    model = build_model(MAX_SEQUENCE_LENGTH + 2)  # We sum 2 for [CLS], [SEP] tokens
    model.load_weights(MODEL_WEIGHTS_DIR)
    model._make_predict_function()