NER: Building the Model

Source

Beginning

Here we'll actually build the model.

  • Feed the data into an Embedding layer, to produce more semantic entries
  • Feed it into an LSTM layer
  • Run the output through a linear layer
  • Run the result through a log softmax layer to get the predicted class for each word.

Imports

# pypi
from trax import layers

# this project
from neurotic.nlp.named_entity_recognition import DataGenerator, NERData, TOKEN

Set Up

ner = NERData()

vocab = vocabulary = ner.data.vocabulary
tag_map = tags = ner.data.tags

Middle

These are the Trax components we'll use (the links are to the implementations on Github).

  • tl.Serial: Combinator that applies layers serially (by function composition).
  • tl.Embedding: Initializes the embedding. In this case it is the dimension of the model by the size of the vocabulary.
    • tl.Embedding(vocab_size, d_feature).
    • vocab_size is the number of unique words in the given vocabulary.
    • d_feature is the number of elements in the word embedding (some choices for a word embedding size range from 150 to 300, for example).
  • tl.LSTM:=Trax= LSTM layer of size d_model.
    • LSTM(n_units) Builds an LSTM layer of n_cells.
  • tl.Dense: A dense layer.
    • tl.Dense(n_units): The parameter n_units is the number of units chosen for this dense layer.
  • tl.LogSoftmax: Log of the output probabilities.
    • Here, you don't need to set any parameters for LogSoftMax().

Online documentation

def NER(vocab_size: int=35181, d_model: int=50, tags: dict=tag_map):
    """
    Args: 
      vocab_size: number of words in the vocabulary
      d_model: the embedding size

    Returns:
       model: a trax serial model
    """
    model = layers.Serial(
        layers.Embedding(vocab_size, d_feature=d_model),
        layers.LSTM(d_model),
        layers.Dense(n_units=len(tag_map)),
        layers.LogSoftmax()
      )
    return model

Inspecting the Model

model = NER()
# display your model
print(model)

Serial[
  Embedding_35181_50
  LSTM_50
  Dense_18
  LogSoftmax
]

Pack It Up for Later

Imports

# python
from collections import namedtuple

# pypi
from trax import layers

import attr

Constants

Settings = namedtuple("Settings", ["embeddings_size"])
SETTINGS = Settings(50)

The Model

@attr.s(auto_attribs=True)
class NER:
    """The named entity recognition model

    Args:
     vocabulary_size: number of tokens in the vocabulary
     tag_count: number of tags
     embeddings_size: the number of features in the embeddings layer
    """
    vocabulary_size: int
    tag_count: int
    embeddings_size: int=SETTINGS.embeddings_size
    _model: layers.Serial=None
  • The Actual Model
    @property
def model(self) -> layers.Serial:
    """The NER model instance"""
    if self._model is None:
        self._model = layers.Serial(
            layers.Embedding(self.vocabulary_size,
                             d_feature=self.embeddings_size),
            layers.LSTM(self.embeddings_size),
            layers.Dense(n_units=self.tag_count),
            layers.LogSoftmax()
      )
    return self._model

Sanity Check

from neurotic.nlp.named_entity_recognition import NER

builder = NER(122, 666)

print(builder.model)

Serial[
  Embedding_122_50
  LSTM_50
  Dense_666
  LogSoftmax
]