From af125875cf89de2a43122c1399b1bb663895034d Mon Sep 17 00:00:00 2001
From: Adriane Boyd <adrianeboyd@gmail.com>
Date: Wed, 5 Aug 2020 14:43:29 +0200
Subject: [PATCH] Update SimpleNER (#5878)

* Fix `get_loss` to use NER annotation
* Add labels as part of cfg
* Add simple overfitting test
---
 spacy/pipeline/simple_ner.py            |  22 +-
 spacy/tests/pipeline/test_simple_ner.py | 449 ++----------------------
 2 files changed, 55 insertions(+), 416 deletions(-)

diff --git a/spacy/pipeline/simple_ner.py b/spacy/pipeline/simple_ner.py
index 44e1182c1..43a3283ca 100644
--- a/spacy/pipeline/simple_ner.py
+++ b/spacy/pipeline/simple_ner.py
@@ -4,12 +4,12 @@ from thinc.api import SequenceCategoricalCrossentropy, set_dropout_rate, Model
 from thinc.api import Optimizer, Config
 from thinc.util import to_numpy
 
+from ..errors import Errors
 from ..gold import Example, spans_from_biluo_tags, iob_to_biluo, biluo_to_iob
 from ..tokens import Doc
 from ..language import Language
 from ..vocab import Vocab
 from ..scorer import Scorer
-from .. import util
 from .pipe import Pipe
 
 
@@ -37,7 +37,6 @@ DEFAULT_SIMPLE_NER_MODEL = Config().from_str(default_model_config)["model"]
     default_config={"labels": [], "model": DEFAULT_SIMPLE_NER_MODEL},
     scores=["ents_p", "ents_r", "ents_f", "ents_per_type"],
     default_score_weights={"ents_f": 1.0, "ents_p": 0.0, "ents_r": 0.0},
-
 )
 def make_simple_ner(
     nlp: Language, name: str, model: Model, labels: Iterable[str]
@@ -60,7 +59,9 @@ class SimpleNER(Pipe):
         self.vocab = vocab
         self.model = model
         self.name = name
-        self.labels = labels
+        self.cfg = {"labels": []}
+        for label in labels:
+            self.add_label(label)
         self.loss_func = SequenceCategoricalCrossentropy(
             names=self.get_tag_names(), normalize=True, missing_value=None
         )
@@ -70,9 +71,20 @@ class SimpleNER(Pipe):
     def is_biluo(self) -> bool:
         return self.model.name.startswith("biluo")
 
+    @property
+    def labels(self) -> Tuple[str]:
+        return tuple(self.cfg["labels"])
+
     def add_label(self, label: str) -> None:
+        """Add a new label to the pipe.
+        label (str): The label to add.
+        DOCS: https://spacy.io/api/simplener#add_label
+        """
+        if not isinstance(label, str):
+            raise ValueError(Errors.E187)
         if label not in self.labels:
-            self.labels.append(label)
+            self.cfg["labels"].append(label)
+            self.vocab.strings.add(label)
 
     def get_tag_names(self) -> List[str]:
         if self.is_biluo:
@@ -133,7 +145,7 @@ class SimpleNER(Pipe):
     def get_loss(self, examples: List[Example], scores) -> Tuple[List[Floats2d], float]:
         truths = []
         for eg in examples:
-            tags = eg.get_aligned("TAG", as_string=True)
+            tags = eg.get_aligned_ner()
             gold_tags = [(tag if tag != "-" else None) for tag in tags]
             if not self.is_biluo:
                 gold_tags = biluo_to_iob(gold_tags)
diff --git a/spacy/tests/pipeline/test_simple_ner.py b/spacy/tests/pipeline/test_simple_ner.py
index 024d7bd26..b012a2cd6 100644
--- a/spacy/tests/pipeline/test_simple_ner.py
+++ b/spacy/tests/pipeline/test_simple_ner.py
@@ -1,418 +1,45 @@
-import pytest
-from collections import namedtuple
-from thinc.api import NumpyOps
-from spacy.ml._biluo import BILUO, _get_transition_table
+from spacy.lang.en import English
+from spacy.gold import Example
+from spacy import util
+from ..util import make_tempdir
 
 
-@pytest.fixture(
-    params=[
-        ["PER", "ORG", "LOC", "MISC"],
-        ["GPE", "PERSON", "NUMBER", "CURRENCY", "EVENT"],
-    ]
-)
-def labels(request):
-    return request.param
+TRAIN_DATA = [
+    ("Who is Shaka Khan?", {"entities": [(7, 17, "PERSON")]}),
+    ("I like London and Berlin.", {"entities": [(7, 13, "LOC"), (18, 24, "LOC")]}),
+]
 
 
-@pytest.fixture
-def ops():
-    return NumpyOps()
+def test_overfitting_IO():
+    # Simple test to try and quickly overfit the SimpleNER component - ensuring the ML models work correctly
+    nlp = English()
+    ner = nlp.add_pipe("simple_ner")
+    train_examples = []
+    for text, annotations in TRAIN_DATA:
+        train_examples.append(Example.from_dict(nlp.make_doc(text), annotations))
+        for ent in annotations.get("entities"):
+            ner.add_label(ent[2])
+    optimizer = nlp.begin_training()
 
+    for i in range(50):
+        losses = {}
+        nlp.update(train_examples, sgd=optimizer, losses=losses)
+    assert losses["ner"] < 0.0001
 
-def _get_actions(labels):
-    action_names = (
-        [f"B{label}" for label in labels]
-        + [f"I{label}" for label in labels]
-        + [f"L{label}" for label in labels]
-        + [f"U{label}" for label in labels]
-        + ["O"]
-    )
-    A = namedtuple("actions", action_names)
-    return A(**{name: i for i, name in enumerate(action_names)})
+    # test the trained model
+    test_text = "I like London."
+    doc = nlp(test_text)
+    ents = doc.ents
+    assert len(ents) == 1
+    assert ents[0].text == "London"
+    assert ents[0].label_ == "LOC"
 
-
-def test_init_biluo_layer(labels):
-    model = BILUO()
-    model.set_dim("nO", model.attrs["get_num_actions"](len(labels)))
-    model.initialize()
-    assert model.get_dim("nO") == len(labels) * 4 + 1
-
-
-def test_transition_table(ops):
-    labels = ["per", "loc", "org"]
-    table = _get_transition_table(len(labels))
-    a = _get_actions(labels)
-    assert table.shape == (2, len(a), len(a))
-    # Not last token, prev action was B
-    assert table[0, a.Bper, a.Bper] == 0
-    assert table[0, a.Bper, a.Bloc] == 0
-    assert table[0, a.Bper, a.Borg] == 0
-    assert table[0, a.Bper, a.Iper] == 1
-    assert table[0, a.Bper, a.Iloc] == 0
-    assert table[0, a.Bper, a.Iorg] == 0
-    assert table[0, a.Bper, a.Lper] == 1
-    assert table[0, a.Bper, a.Lloc] == 0
-    assert table[0, a.Bper, a.Lorg] == 0
-    assert table[0, a.Bper, a.Uper] == 0
-    assert table[0, a.Bper, a.Uloc] == 0
-    assert table[0, a.Bper, a.Uorg] == 0
-    assert table[0, a.Bper, a.O] == 0
-
-    assert table[0, a.Bloc, a.Bper] == 0
-    assert table[0, a.Bloc, a.Bloc] == 0
-    assert table[0, a.Bloc, a.Borg] == 0
-    assert table[0, a.Bloc, a.Iper] == 0
-    assert table[0, a.Bloc, a.Iloc] == 1
-    assert table[0, a.Bloc, a.Iorg] == 0
-    assert table[0, a.Bloc, a.Lper] == 0
-    assert table[0, a.Bloc, a.Lloc] == 1
-    assert table[0, a.Bloc, a.Lorg] == 0
-    assert table[0, a.Bloc, a.Uper] == 0
-    assert table[0, a.Bloc, a.Uloc] == 0
-    assert table[0, a.Bloc, a.Uorg] == 0
-    assert table[0, a.Bloc, a.O] == 0
-
-    assert table[0, a.Borg, a.Bper] == 0
-    assert table[0, a.Borg, a.Bloc] == 0
-    assert table[0, a.Borg, a.Borg] == 0
-    assert table[0, a.Borg, a.Iper] == 0
-    assert table[0, a.Borg, a.Iloc] == 0
-    assert table[0, a.Borg, a.Iorg] == 1
-    assert table[0, a.Borg, a.Lper] == 0
-    assert table[0, a.Borg, a.Lloc] == 0
-    assert table[0, a.Borg, a.Lorg] == 1
-    assert table[0, a.Borg, a.Uper] == 0
-    assert table[0, a.Borg, a.Uloc] == 0
-    assert table[0, a.Borg, a.Uorg] == 0
-    assert table[0, a.Borg, a.O] == 0
-
-    # Not last token, prev action was I
-    assert table[0, a.Iper, a.Bper] == 0
-    assert table[0, a.Iper, a.Bloc] == 0
-    assert table[0, a.Iper, a.Borg] == 0
-    assert table[0, a.Iper, a.Iper] == 1
-    assert table[0, a.Iper, a.Iloc] == 0
-    assert table[0, a.Iper, a.Iorg] == 0
-    assert table[0, a.Iper, a.Lper] == 1
-    assert table[0, a.Iper, a.Lloc] == 0
-    assert table[0, a.Iper, a.Lorg] == 0
-    assert table[0, a.Iper, a.Uper] == 0
-    assert table[0, a.Iper, a.Uloc] == 0
-    assert table[0, a.Iper, a.Uorg] == 0
-    assert table[0, a.Iper, a.O] == 0
-
-    assert table[0, a.Iloc, a.Bper] == 0
-    assert table[0, a.Iloc, a.Bloc] == 0
-    assert table[0, a.Iloc, a.Borg] == 0
-    assert table[0, a.Iloc, a.Iper] == 0
-    assert table[0, a.Iloc, a.Iloc] == 1
-    assert table[0, a.Iloc, a.Iorg] == 0
-    assert table[0, a.Iloc, a.Lper] == 0
-    assert table[0, a.Iloc, a.Lloc] == 1
-    assert table[0, a.Iloc, a.Lorg] == 0
-    assert table[0, a.Iloc, a.Uper] == 0
-    assert table[0, a.Iloc, a.Uloc] == 0
-    assert table[0, a.Iloc, a.Uorg] == 0
-    assert table[0, a.Iloc, a.O] == 0
-
-    assert table[0, a.Iorg, a.Bper] == 0
-    assert table[0, a.Iorg, a.Bloc] == 0
-    assert table[0, a.Iorg, a.Borg] == 0
-    assert table[0, a.Iorg, a.Iper] == 0
-    assert table[0, a.Iorg, a.Iloc] == 0
-    assert table[0, a.Iorg, a.Iorg] == 1
-    assert table[0, a.Iorg, a.Lper] == 0
-    assert table[0, a.Iorg, a.Lloc] == 0
-    assert table[0, a.Iorg, a.Lorg] == 1
-    assert table[0, a.Iorg, a.Uper] == 0
-    assert table[0, a.Iorg, a.Uloc] == 0
-    assert table[0, a.Iorg, a.Uorg] == 0
-    assert table[0, a.Iorg, a.O] == 0
-
-    # Not last token, prev action was L
-    assert table[0, a.Lper, a.Bper] == 1
-    assert table[0, a.Lper, a.Bloc] == 1
-    assert table[0, a.Lper, a.Borg] == 1
-    assert table[0, a.Lper, a.Iper] == 0
-    assert table[0, a.Lper, a.Iloc] == 0
-    assert table[0, a.Lper, a.Iorg] == 0
-    assert table[0, a.Lper, a.Lper] == 0
-    assert table[0, a.Lper, a.Lloc] == 0
-    assert table[0, a.Lper, a.Lorg] == 0
-    assert table[0, a.Lper, a.Uper] == 1
-    assert table[0, a.Lper, a.Uloc] == 1
-    assert table[0, a.Lper, a.Uorg] == 1
-    assert table[0, a.Lper, a.O] == 1
-
-    assert table[0, a.Lloc, a.Bper] == 1
-    assert table[0, a.Lloc, a.Bloc] == 1
-    assert table[0, a.Lloc, a.Borg] == 1
-    assert table[0, a.Lloc, a.Iper] == 0
-    assert table[0, a.Lloc, a.Iloc] == 0
-    assert table[0, a.Lloc, a.Iorg] == 0
-    assert table[0, a.Lloc, a.Lper] == 0
-    assert table[0, a.Lloc, a.Lloc] == 0
-    assert table[0, a.Lloc, a.Lorg] == 0
-    assert table[0, a.Lloc, a.Uper] == 1
-    assert table[0, a.Lloc, a.Uloc] == 1
-    assert table[0, a.Lloc, a.Uorg] == 1
-    assert table[0, a.Lloc, a.O] == 1
-
-    assert table[0, a.Lorg, a.Bper] == 1
-    assert table[0, a.Lorg, a.Bloc] == 1
-    assert table[0, a.Lorg, a.Borg] == 1
-    assert table[0, a.Lorg, a.Iper] == 0
-    assert table[0, a.Lorg, a.Iloc] == 0
-    assert table[0, a.Lorg, a.Iorg] == 0
-    assert table[0, a.Lorg, a.Lper] == 0
-    assert table[0, a.Lorg, a.Lloc] == 0
-    assert table[0, a.Lorg, a.Lorg] == 0
-    assert table[0, a.Lorg, a.Uper] == 1
-    assert table[0, a.Lorg, a.Uloc] == 1
-    assert table[0, a.Lorg, a.Uorg] == 1
-    assert table[0, a.Lorg, a.O] == 1
-
-    # Not last token, prev action was U
-    assert table[0, a.Uper, a.Bper] == 1
-    assert table[0, a.Uper, a.Bloc] == 1
-    assert table[0, a.Uper, a.Borg] == 1
-    assert table[0, a.Uper, a.Iper] == 0
-    assert table[0, a.Uper, a.Iloc] == 0
-    assert table[0, a.Uper, a.Iorg] == 0
-    assert table[0, a.Uper, a.Lper] == 0
-    assert table[0, a.Uper, a.Lloc] == 0
-    assert table[0, a.Uper, a.Lorg] == 0
-    assert table[0, a.Uper, a.Uper] == 1
-    assert table[0, a.Uper, a.Uloc] == 1
-    assert table[0, a.Uper, a.Uorg] == 1
-    assert table[0, a.Uper, a.O] == 1
-
-    assert table[0, a.Uloc, a.Bper] == 1
-    assert table[0, a.Uloc, a.Bloc] == 1
-    assert table[0, a.Uloc, a.Borg] == 1
-    assert table[0, a.Uloc, a.Iper] == 0
-    assert table[0, a.Uloc, a.Iloc] == 0
-    assert table[0, a.Uloc, a.Iorg] == 0
-    assert table[0, a.Uloc, a.Lper] == 0
-    assert table[0, a.Uloc, a.Lloc] == 0
-    assert table[0, a.Uloc, a.Lorg] == 0
-    assert table[0, a.Uloc, a.Uper] == 1
-    assert table[0, a.Uloc, a.Uloc] == 1
-    assert table[0, a.Uloc, a.Uorg] == 1
-    assert table[0, a.Uloc, a.O] == 1
-
-    assert table[0, a.Uorg, a.Bper] == 1
-    assert table[0, a.Uorg, a.Bloc] == 1
-    assert table[0, a.Uorg, a.Borg] == 1
-    assert table[0, a.Uorg, a.Iper] == 0
-    assert table[0, a.Uorg, a.Iloc] == 0
-    assert table[0, a.Uorg, a.Iorg] == 0
-    assert table[0, a.Uorg, a.Lper] == 0
-    assert table[0, a.Uorg, a.Lloc] == 0
-    assert table[0, a.Uorg, a.Lorg] == 0
-    assert table[0, a.Uorg, a.Uper] == 1
-    assert table[0, a.Uorg, a.Uloc] == 1
-    assert table[0, a.Uorg, a.Uorg] == 1
-    assert table[0, a.Uorg, a.O] == 1
-
-    # Not last token, prev action was O
-    assert table[0, a.O, a.Bper] == 1
-    assert table[0, a.O, a.Bloc] == 1
-    assert table[0, a.O, a.Borg] == 1
-    assert table[0, a.O, a.Iper] == 0
-    assert table[0, a.O, a.Iloc] == 0
-    assert table[0, a.O, a.Iorg] == 0
-    assert table[0, a.O, a.Lper] == 0
-    assert table[0, a.O, a.Lloc] == 0
-    assert table[0, a.O, a.Lorg] == 0
-    assert table[0, a.O, a.Uper] == 1
-    assert table[0, a.O, a.Uloc] == 1
-    assert table[0, a.O, a.Uorg] == 1
-    assert table[0, a.O, a.O] == 1
-
-    # Last token, prev action was B
-    assert table[1, a.Bper, a.Bper] == 0
-    assert table[1, a.Bper, a.Bloc] == 0
-    assert table[1, a.Bper, a.Borg] == 0
-    assert table[1, a.Bper, a.Iper] == 0
-    assert table[1, a.Bper, a.Iloc] == 0
-    assert table[1, a.Bper, a.Iorg] == 0
-    assert table[1, a.Bper, a.Lper] == 1
-    assert table[1, a.Bper, a.Lloc] == 0
-    assert table[1, a.Bper, a.Lorg] == 0
-    assert table[1, a.Bper, a.Uper] == 0
-    assert table[1, a.Bper, a.Uloc] == 0
-    assert table[1, a.Bper, a.Uorg] == 0
-    assert table[1, a.Bper, a.O] == 0
-
-    assert table[1, a.Bloc, a.Bper] == 0
-    assert table[1, a.Bloc, a.Bloc] == 0
-    assert table[0, a.Bloc, a.Borg] == 0
-    assert table[1, a.Bloc, a.Iper] == 0
-    assert table[1, a.Bloc, a.Iloc] == 0
-    assert table[1, a.Bloc, a.Iorg] == 0
-    assert table[1, a.Bloc, a.Lper] == 0
-    assert table[1, a.Bloc, a.Lloc] == 1
-    assert table[1, a.Bloc, a.Lorg] == 0
-    assert table[1, a.Bloc, a.Uper] == 0
-    assert table[1, a.Bloc, a.Uloc] == 0
-    assert table[1, a.Bloc, a.Uorg] == 0
-    assert table[1, a.Bloc, a.O] == 0
-
-    assert table[1, a.Borg, a.Bper] == 0
-    assert table[1, a.Borg, a.Bloc] == 0
-    assert table[1, a.Borg, a.Borg] == 0
-    assert table[1, a.Borg, a.Iper] == 0
-    assert table[1, a.Borg, a.Iloc] == 0
-    assert table[1, a.Borg, a.Iorg] == 0
-    assert table[1, a.Borg, a.Lper] == 0
-    assert table[1, a.Borg, a.Lloc] == 0
-    assert table[1, a.Borg, a.Lorg] == 1
-    assert table[1, a.Borg, a.Uper] == 0
-    assert table[1, a.Borg, a.Uloc] == 0
-    assert table[1, a.Borg, a.Uorg] == 0
-    assert table[1, a.Borg, a.O] == 0
-
-    # Last token, prev action was I
-    assert table[1, a.Iper, a.Bper] == 0
-    assert table[1, a.Iper, a.Bloc] == 0
-    assert table[1, a.Iper, a.Borg] == 0
-    assert table[1, a.Iper, a.Iper] == 0
-    assert table[1, a.Iper, a.Iloc] == 0
-    assert table[1, a.Iper, a.Iorg] == 0
-    assert table[1, a.Iper, a.Lper] == 1
-    assert table[1, a.Iper, a.Lloc] == 0
-    assert table[1, a.Iper, a.Lorg] == 0
-    assert table[1, a.Iper, a.Uper] == 0
-    assert table[1, a.Iper, a.Uloc] == 0
-    assert table[1, a.Iper, a.Uorg] == 0
-    assert table[1, a.Iper, a.O] == 0
-
-    assert table[1, a.Iloc, a.Bper] == 0
-    assert table[1, a.Iloc, a.Bloc] == 0
-    assert table[1, a.Iloc, a.Borg] == 0
-    assert table[1, a.Iloc, a.Iper] == 0
-    assert table[1, a.Iloc, a.Iloc] == 0
-    assert table[1, a.Iloc, a.Iorg] == 0
-    assert table[1, a.Iloc, a.Lper] == 0
-    assert table[1, a.Iloc, a.Lloc] == 1
-    assert table[1, a.Iloc, a.Lorg] == 0
-    assert table[1, a.Iloc, a.Uper] == 0
-    assert table[1, a.Iloc, a.Uloc] == 0
-    assert table[1, a.Iloc, a.Uorg] == 0
-    assert table[1, a.Iloc, a.O] == 0
-
-    assert table[1, a.Iorg, a.Bper] == 0
-    assert table[1, a.Iorg, a.Bloc] == 0
-    assert table[1, a.Iorg, a.Borg] == 0
-    assert table[1, a.Iorg, a.Iper] == 0
-    assert table[1, a.Iorg, a.Iloc] == 0
-    assert table[1, a.Iorg, a.Iorg] == 0
-    assert table[1, a.Iorg, a.Lper] == 0
-    assert table[1, a.Iorg, a.Lloc] == 0
-    assert table[1, a.Iorg, a.Lorg] == 1
-    assert table[1, a.Iorg, a.Uper] == 0
-    assert table[1, a.Iorg, a.Uloc] == 0
-    assert table[1, a.Iorg, a.Uorg] == 0
-    assert table[1, a.Iorg, a.O] == 0
-
-    # Last token, prev action was L
-    assert table[1, a.Lper, a.Bper] == 0
-    assert table[1, a.Lper, a.Bloc] == 0
-    assert table[1, a.Lper, a.Borg] == 0
-    assert table[1, a.Lper, a.Iper] == 0
-    assert table[1, a.Lper, a.Iloc] == 0
-    assert table[1, a.Lper, a.Iorg] == 0
-    assert table[1, a.Lper, a.Lper] == 0
-    assert table[1, a.Lper, a.Lloc] == 0
-    assert table[1, a.Lper, a.Lorg] == 0
-    assert table[1, a.Lper, a.Uper] == 1
-    assert table[1, a.Lper, a.Uloc] == 1
-    assert table[1, a.Lper, a.Uorg] == 1
-    assert table[1, a.Lper, a.O] == 1
-
-    assert table[1, a.Lloc, a.Bper] == 0
-    assert table[1, a.Lloc, a.Bloc] == 0
-    assert table[1, a.Lloc, a.Borg] == 0
-    assert table[1, a.Lloc, a.Iper] == 0
-    assert table[1, a.Lloc, a.Iloc] == 0
-    assert table[1, a.Lloc, a.Iorg] == 0
-    assert table[1, a.Lloc, a.Lper] == 0
-    assert table[1, a.Lloc, a.Lloc] == 0
-    assert table[1, a.Lloc, a.Lorg] == 0
-    assert table[1, a.Lloc, a.Uper] == 1
-    assert table[1, a.Lloc, a.Uloc] == 1
-    assert table[1, a.Lloc, a.Uorg] == 1
-    assert table[1, a.Lloc, a.O] == 1
-
-    assert table[1, a.Lorg, a.Bper] == 0
-    assert table[1, a.Lorg, a.Bloc] == 0
-    assert table[1, a.Lorg, a.Borg] == 0
-    assert table[1, a.Lorg, a.Iper] == 0
-    assert table[1, a.Lorg, a.Iloc] == 0
-    assert table[1, a.Lorg, a.Iorg] == 0
-    assert table[1, a.Lorg, a.Lper] == 0
-    assert table[1, a.Lorg, a.Lloc] == 0
-    assert table[1, a.Lorg, a.Lorg] == 0
-    assert table[1, a.Lorg, a.Uper] == 1
-    assert table[1, a.Lorg, a.Uloc] == 1
-    assert table[1, a.Lorg, a.Uorg] == 1
-    assert table[1, a.Lorg, a.O] == 1
-
-    # Last token, prev action was U
-    assert table[1, a.Uper, a.Bper] == 0
-    assert table[1, a.Uper, a.Bloc] == 0
-    assert table[1, a.Uper, a.Borg] == 0
-    assert table[1, a.Uper, a.Iper] == 0
-    assert table[1, a.Uper, a.Iloc] == 0
-    assert table[1, a.Uper, a.Iorg] == 0
-    assert table[1, a.Uper, a.Lper] == 0
-    assert table[1, a.Uper, a.Lloc] == 0
-    assert table[1, a.Uper, a.Lorg] == 0
-    assert table[1, a.Uper, a.Uper] == 1
-    assert table[1, a.Uper, a.Uloc] == 1
-    assert table[1, a.Uper, a.Uorg] == 1
-    assert table[1, a.Uper, a.O] == 1
-
-    assert table[1, a.Uloc, a.Bper] == 0
-    assert table[1, a.Uloc, a.Bloc] == 0
-    assert table[1, a.Uloc, a.Borg] == 0
-    assert table[1, a.Uloc, a.Iper] == 0
-    assert table[1, a.Uloc, a.Iloc] == 0
-    assert table[1, a.Uloc, a.Iorg] == 0
-    assert table[1, a.Uloc, a.Lper] == 0
-    assert table[1, a.Uloc, a.Lloc] == 0
-    assert table[1, a.Uloc, a.Lorg] == 0
-    assert table[1, a.Uloc, a.Uper] == 1
-    assert table[1, a.Uloc, a.Uloc] == 1
-    assert table[1, a.Uloc, a.Uorg] == 1
-    assert table[1, a.Uloc, a.O] == 1
-
-    assert table[1, a.Uorg, a.Bper] == 0
-    assert table[1, a.Uorg, a.Bloc] == 0
-    assert table[1, a.Uorg, a.Borg] == 0
-    assert table[1, a.Uorg, a.Iper] == 0
-    assert table[1, a.Uorg, a.Iloc] == 0
-    assert table[1, a.Uorg, a.Iorg] == 0
-    assert table[1, a.Uorg, a.Lper] == 0
-    assert table[1, a.Uorg, a.Lloc] == 0
-    assert table[1, a.Uorg, a.Lorg] == 0
-    assert table[1, a.Uorg, a.Uper] == 1
-    assert table[1, a.Uorg, a.Uloc] == 1
-    assert table[1, a.Uorg, a.Uorg] == 1
-    assert table[1, a.Uorg, a.O] == 1
-
-    # Last token, prev action was O
-    assert table[1, a.O, a.Bper] == 0
-    assert table[1, a.O, a.Bloc] == 0
-    assert table[1, a.O, a.Borg] == 0
-    assert table[1, a.O, a.Iper] == 0
-    assert table[1, a.O, a.Iloc] == 0
-    assert table[1, a.O, a.Iorg] == 0
-    assert table[1, a.O, a.Lper] == 0
-    assert table[1, a.O, a.Lloc] == 0
-    assert table[1, a.O, a.Lorg] == 0
-    assert table[1, a.O, a.Uper] == 1
-    assert table[1, a.O, a.Uloc] == 1
-    assert table[1, a.O, a.Uorg] == 1
-    assert table[1, a.O, a.O] == 1
+    # Also test the results are still the same after IO
+    with make_tempdir() as tmp_dir:
+        nlp.to_disk(tmp_dir)
+        nlp2 = util.load_model_from_path(tmp_dir)
+        doc2 = nlp2(test_text)
+        ents2 = doc2.ents
+        assert len(ents2) == 1
+        assert ents2[0].text == "London"
+        assert ents2[0].label_ == "LOC"