history.py

history.py

def override_loss_hook(ner_loss):
    old_loss_hook = ner_loss.loss_hook

    history = dict(preds=[])

    def loss_hook(self, pred, target):
        history['preds'].append(pred)
        return {}

    ner_loss.loss_hook = types.MethodType(loss_hook, ner_loss)

    def reset_func():
        ner_loss.loss_hook = old_loss_hook

    return history, reset_func


if __name__ == '__main__':
    ner_loss = NERLoss()
    history, reset_loss_hook = override_loss_hook(ner_loss)

    for batch in batches:
        ner_loss(batch) # this will also append the pred to history

    aggregate(history)

    reset_loss_hook()