NER: Training the Model

# from python
from collections import namedtuple
from functools import partial
from tempfile import TemporaryFile

import random
import sys

# from pypi
from holoviews import opts
from trax import layers
from trax.supervised import training

import holoviews
import hvplot.pandas
import pandas
import trax

# this project
from neurotic.nlp.named_entity_recognition import (DataGenerator,
                                                   NER,
                                                   NERData,
                                                   TOKEN)
# another project
from graeae import EmbedHoloviews, Timer

train_generator = trax.data.inputs.add_loss_weights(
    training_generator,
    id_to_mask=ner.data.vocabulary[TOKEN.pad])

evaluate_generator = trax.data.inputs.add_loss_weights(
    validation_generator,
    id_to_mask=ner.data.vocabulary[TOKEN.pad])

def train_model(NER: trax.layers.Serial,
                train_generator: type,
                eval_generator: type,
                train_steps: int=1,
                steps_per_checkpoint: int=100,
                learning_rate: float=SETTINGS.learning_rate,
                verbose: bool=False,
                output_dir="~/models/ner/") -> training.Loop:
    """Train the Named Entity Recognition Model
    Args: 
      NER: the model you are building
      train_generator: The data generator for training examples
      eval_generator: The data generator for validation examples,
      train_steps: number of training steps
      output_dir: folder to save your model

    Returns:
      training_loop: a trax supervised training Loop
    """
    train_task = training.TrainTask(
        labeled_data=train_generator,
        loss_layer = layers.WeightedCategoryCrossEntropy(),
        optimizer = trax.optimizers.Adam(learning_rate),
        n_steps_per_checkpoint=steps_per_checkpoint,
    )

    eval_task = training.EvalTask(
      labeled_data = eval_generator,
      metrics = [layers.WeightedCategoryCrossEntropy(),
                 layers.Accuracy()],
      n_eval_batches = SETTINGS.batch_size
    )

    training_loop = training.Loop(
        NER,
        train_task,
        eval_tasks=[eval_task],
        output_dir=output_dir)

    if verbose:
        print(f"Running {train_steps} steps")
    training_loop.run(n_steps = train_steps)
    return training_loop

training_steps = 1500
real_stdout = sys.stdout

TIMER.emit = False
TIMER.start()
with TemporaryFile("w") as temp_file:
    sys.stdout = temp_file
    training_loop = train_model(trainee.model, train_generator,
                                evaluate_generator,
                                steps_per_checkpoint=10,
                                train_steps=training_steps,
                                verbose=False)
TIMER.stop()
sys.stdout = real_stdout

print(f"{TIMER.ended - TIMER.started}")