Fix replacing Transformer listener (#277)

* failing test for replace listeners

* replace listener attr

* add methods for IO of inhouse Tok2VecTransformer

* add custom hack to rewrite transformer config

* cleanup

* to_bytes and from_bytes fails for transformer-based pipeline

* simple from_bytes failing test

* from_bytes works if we call initialize first

* fix test

* ensure predictions are still the same after replacing the listener

* clean up comments

* bump versions

* fix initialize nlp default value

* add comment

Author: svlandeg

requirements.txt

@@ -1,4 +1,4 @@
-spacy>=3.0.0,<4.0.0
+spacy>=3.1.0,<4.0.0
 transformers>=3.4.0,<4.10.0
 torch>=1.5.0
 srsly>=2.4.0,<3.0.0

setup.cfg

@@ -1,5 +1,5 @@
 [metadata]
-version = 1.0.4
+version = 1.0.3
 description = spaCy pipelines for pre-trained BERT and other transformers
 url = https://spacy.io
 author = Explosion
@@ -27,7 +27,7 @@ zip_safe = false
 include_package_data = true
 python_requires = >=3.6
 install_requires =
-    spacy>=3.0.0,<4.0.0
+    spacy>=3.1.0,<4.0.0
     transformers>=3.4.0,<4.10.0
     torch>=1.5.0
     srsly>=2.4.0,<3.0.0

spacy_transformers/architectures.py

@@ -42,6 +42,7 @@ def transformer_listener_tok2vec_v1(
     return model
 
 
+# Note: when updating, also make sure to update 'replace_listener_cfg' in _util.py
 @registry.architectures.register("spacy-transformers.Tok2VecTransformer.v1")
 def transformer_tok2vec_v1(
     name: str,

spacy_transformers/layers/_util.py

@@ -0,0 +1,21 @@
+from thinc.api import Model, chain
+from .split_trf import split_trf_batch
+
+
+def replace_listener(model):
+    return chain(model, split_trf_batch())
+
+
+def replace_listener_cfg(tok2vec_model_cfg, listener_model_cfg):
+    result = tok2vec_model_cfg
+    if (
+        "TransformerModel" in tok2vec_model_cfg["@architectures"]
+        and "TransformerListener" in listener_model_cfg["@architectures"]
+    ):
+        result["@architectures"] = "spacy-transformers.Tok2VecTransformer.v1"
+        if "pooling" in listener_model_cfg and "pooling" not in result:
+            result["pooling"] = listener_model_cfg["pooling"]
+        if "grad_factor" in listener_model_cfg and "grad_factor" not in result:
+            result["grad_factor"] = listener_model_cfg["pooling"]
+
+    return result

spacy_transformers/layers/split_trf.py

@@ -1,4 +1,4 @@
-from thinc.api import Model
+from thinc.api import Model, chain
 from typing import List
 from ..data_classes import FullTransformerBatch, TransformerData
 

spacy_transformers/layers/transformer_model.py

@@ -1,16 +1,17 @@
 from typing import List, Tuple, Callable
+
+from spacy_transformers.layers._util import replace_listener, replace_listener_cfg
+from thinc.types import ArgsKwargs
 import torch
 from spacy.tokens import Doc
-from thinc.api import PyTorchWrapper, Model, xp2torch
-from thinc.types import ArgsKwargs
-from transformers.tokenization_utils import BatchEncoding
+from thinc.api import PyTorchWrapper, Model, xp2torch, chain
 import logging
 
 from ..data_classes import FullTransformerBatch, WordpieceBatch
 from ..util import huggingface_tokenize, huggingface_from_pretrained
 from ..util import find_last_hidden, maybe_flush_pytorch_cache
-from ..truncate import truncate_oversize_splits
 from ..util import log_gpu_memory, log_batch_size
+from ..truncate import truncate_oversize_splits
 from ..align import get_alignment
 
 
@@ -41,6 +42,8 @@ def TransformerModel(
             "set_transformer": set_pytorch_transformer,
             "has_transformer": False,
             "flush_cache_chance": 0.0,
+            "replace_listener": replace_listener,
+            "replace_listener_cfg": replace_listener_cfg,
         },
     )
 

spacy_transformers/pipeline_component.py

@@ -328,7 +328,7 @@ def get_loss(self, docs, golds, scores):
         pass
 
     def initialize(
-        self, get_examples: Callable[[], Iterable[Example]], *, nlp: Optional[Language]
+        self, get_examples: Callable[[], Iterable[Example]], *, nlp: Optional[Language] = None
     ):
         """Initialize the pipe for training, using data examples if available.
 

spacy_transformers/tests/test_pipeline_component.py

@@ -145,7 +145,7 @@ def test_transformer_pipeline_long_token(simple_nlp):
     """
 
 
-def test_transformer_pipeline_tagger():
+def test_transformer_pipeline_tagger_listener():
     """Test that a pipeline with just a transformer+tagger runs and trains properly"""
     orig_config = Config().from_str(cfg_string)
     nlp = util.load_model_from_config(orig_config, auto_fill=True, validate=True)
@@ -171,7 +171,8 @@ def test_transformer_pipeline_tagger():
         losses = {}
         nlp.update(train_examples, sgd=optimizer, losses=losses)
 
-    doc = nlp("We're interested at underwater basket weaving.")
+    text = "We're interested at underwater basket weaving."
+    doc = nlp(text)
     doc_tensor = tagger_trf.predict([doc])
     assert_equal(doc._.trf_data.tensors, doc_tensor[0].tensors)
 
@@ -180,12 +181,23 @@ def test_transformer_pipeline_tagger():
         file_path = d / "trained_nlp"
         nlp.to_disk(file_path)
         nlp2 = util.load_model_from_path(file_path)
-        doc = nlp2("We're interested at underwater basket weaving.")
+        doc = nlp2(text)
         tagger2 = nlp2.get_pipe("tagger")
         tagger_trf2 = tagger2.model.get_ref("tok2vec").layers[0]
         doc_tensor2 = tagger_trf2.predict([doc])
         assert_equal(doc_tensor2[0].tensors, doc_tensor[0].tensors)
 
+    # ensure to_bytes / from_bytes works
+    nlp_bytes = nlp.to_bytes()
+    nlp3 = util.load_model_from_config(orig_config, auto_fill=True, validate=True)
+    nlp3.initialize(lambda: train_examples)
+    nlp3.from_bytes(nlp_bytes)
+    doc = nlp3(text)
+    tagger3 = nlp3.get_pipe("tagger")
+    tagger_trf3 = tagger3.model.get_ref("tok2vec").layers[0]
+    doc_tensor3 = tagger_trf3.predict([doc])
+    assert_equal(doc_tensor3[0].tensors, doc_tensor[0].tensors)
+
 
 def test_transformer_pipeline_empty():
     """Test that the pipeline doesn't fail with empty input"""
@@ -255,3 +267,71 @@ def test_multiprocessing(simple_nlp, texts):
 
         for doc, expected_doc in zip(docs, expecteds):
             assert_docs_equal(doc, expected_doc)
+
+def test_replace_listeners():
+    orig_config = Config().from_str(cfg_string)
+    nlp = util.load_model_from_config(orig_config, auto_fill=True, validate=True)
+    text = "This is awesome"
+    examples = [Example.from_dict(nlp.make_doc(text), {"tags": ["A", "B", "C"]})]
+    optimizer = nlp.initialize(lambda: examples)
+    # verify correct configuration with transformer listener
+    transformer = nlp.get_pipe("transformer")
+    tagger = nlp.get_pipe("tagger")
+    tagger_tok2vec = tagger.model.get_ref("tok2vec")
+    tagger_listener = tagger_tok2vec.get_ref("listener")
+    assert isinstance(tagger_listener, TransformerListener)
+    assert transformer.listener_map["tagger"][0] == tagger_listener
+    assert (
+        nlp.config["components"]["transformer"]["model"]["@architectures"]
+        == "spacy-transformers.TransformerModel.v1"
+    )
+    assert (
+        nlp.config["components"]["tagger"]["model"]["tok2vec"]["@architectures"]
+        == "spacy-transformers.TransformerListener.v1"
+    )
+    # train pipe before replacing listeners
+    for i in range(2):
+        losses = {}
+        nlp.update(examples, sgd=optimizer, losses=losses)
+        doc = nlp(text)
+
+    preds = [t.tag_ for t in doc]
+    doc_tensor = tagger_tok2vec.predict([doc])
+
+    # replace listener and verify predictions are still the same
+    nlp.replace_listeners("transformer", "tagger", ["model.tok2vec"])
+    tagger = nlp.get_pipe("tagger")
+    tagger_tok2vec = tagger.model.get_ref("tok2vec")
+    assert tagger_tok2vec.layers[0].layers[0].name == "transformer"
+    assert (
+        nlp.config["components"]["tagger"]["model"]["tok2vec"]["@architectures"]
+        == "spacy-transformers.Tok2VecTransformer.v1"
+    )
+    doc2 = nlp(text)
+    assert preds == [t.tag_ for t in doc2]
+    assert_equal(doc_tensor, tagger_tok2vec.predict([doc2]))
+    # attempt training with the new pipeline
+    optimizer = nlp.resume_training()
+    for i in range(2):
+        losses = {}
+        nlp.update(examples, sgd=optimizer, losses=losses)
+        assert losses["tagger"] > 0.0
+
+
+def test_replace_listeners_invalid():
+    orig_config = Config().from_str(cfg_string)
+    nlp = util.load_model_from_config(orig_config, auto_fill=True, validate=True)
+    text = "This is awesome"
+    examples = [Example.from_dict(nlp.make_doc(text), {"tags": ["A", "B", "C"]})]
+    optimizer = nlp.initialize(lambda: examples)
+    for i in range(2):
+        losses = {}
+        nlp.update(examples, sgd=optimizer, losses=losses)
+    with pytest.raises(ValueError):
+        nlp.replace_listeners("invalid", "tagger", ["model.tok2vec"])
+    with pytest.raises(ValueError):
+        nlp.replace_listeners("transformer", "parser", ["model.tok2vec"])
+    with pytest.raises(ValueError):
+        nlp.replace_listeners("transformer", "tagger", ["model.yolo"])
+    with pytest.raises(ValueError):
+        nlp.replace_listeners("transformer", "tagger", ["model.tok2vec", "model.yolo"])

spacy_transformers/tests/test_serialize.py

@@ -1,3 +1,5 @@
+from spacy import Language
+
 from spacy_transformers import TransformerData
 import srsly
 
@@ -7,3 +9,25 @@ def test_serialize_transformer_data():
     bytes_data = srsly.msgpack_dumps(data)
     new_data = srsly.msgpack_loads(bytes_data)
     assert isinstance(new_data["x"], TransformerData)
+
+
+def test_transformer_tobytes():
+    nlp = Language()
+    trf = nlp.add_pipe("transformer")
+    trf_bytes = trf.to_bytes()
+
+    nlp2 = Language()
+    trf2 = nlp2.add_pipe("transformer")
+    trf2.from_bytes(trf_bytes)
+
+
+def test_transformer_model_tobytes():
+    nlp = Language()
+    trf = nlp.add_pipe("transformer")
+    nlp.initialize()
+    trf_bytes = trf.to_bytes()
+
+    nlp2 = Language()
+    trf2 = nlp2.add_pipe("transformer")
+    nlp2.initialize()
+    trf2.from_bytes(trf_bytes)

spacy_transformers/tests/test_tok2vectransformer.py

@@ -0,0 +1,88 @@
+import pytest
+from spacy.training.example import Example
+from spacy.util import make_tempdir
+from spacy import util
+from thinc.api import Model, Config
+from numpy.testing import assert_equal
+
+
+TRAIN_DATA = [
+    ("I like green eggs", {"tags": ["N", "V", "J", "N"]}),
+    ("Eat blue ham", {"tags": ["V", "J", "N"]}),
+]
+
+
+cfg_string = """
+    [nlp]
+    lang = "en"
+    pipeline = ["tagger"]
+
+    [components]
+
+    [components.tagger]
+    factory = "tagger"
+
+    [components.tagger.model]
+    @architectures = "spacy.Tagger.v1"
+    nO = null
+
+    [components.tagger.model.tok2vec]
+    @architectures = "spacy-transformers.Tok2VecTransformer.v1"
+    name = "albert-base-v2"
+    tokenizer_config = {"use_fast": false}
+    grad_factor = 1.0
+
+    [components.tagger.model.tok2vec.get_spans]
+    @span_getters = "spacy-transformers.strided_spans.v1"
+    window = 256
+    stride = 256
+
+    [components.tagger.model.tok2vec.pooling]
+    @layers = "reduce_mean.v1"
+    """
+
+
+def test_transformer_pipeline_tagger_internal():
+    """Test that a tagger with internal transformer runs and trains properly"""
+    orig_config = Config().from_str(cfg_string)
+    nlp = util.load_model_from_config(orig_config, auto_fill=True, validate=True)
+    assert nlp.pipe_names == ["tagger"]
+    tagger = nlp.get_pipe("tagger")
+    tagger_trf = tagger.model.get_ref("tok2vec").layers[0]
+    assert isinstance(tagger_trf, Model)
+    train_examples = []
+    for t in TRAIN_DATA:
+        train_examples.append(Example.from_dict(nlp.make_doc(t[0]), t[1]))
+        for tag in t[1]["tags"]:
+            tagger.add_label(tag)
+
+    optimizer = nlp.initialize(lambda: train_examples)
+    for i in range(2):
+        losses = {}
+        nlp.update(train_examples, sgd=optimizer, losses=losses)
+
+    doc = nlp("We're interested at underwater basket weaving.")
+    doc_tensor = tagger_trf.predict([doc])
+
+    # ensure IO goes OK
+    with make_tempdir() as d:
+        file_path = d / "trained_nlp"
+        nlp.to_disk(file_path)
+        nlp2 = util.load_model_from_config(orig_config, auto_fill=True, validate=True)
+        nlp2.initialize(lambda: train_examples)
+
+        # results are not the same if we don't call from_disk
+        doc2 = nlp2("We're interested at underwater basket weaving.")
+        tagger2 = nlp2.get_pipe("tagger")
+        tagger_trf2 = tagger2.model.get_ref("tok2vec").layers[0]
+        doc_tensor2 = tagger_trf2.predict([doc2])
+        with pytest.raises(AssertionError):
+            assert_equal(doc_tensor2.doc_data[0].tensors, doc_tensor.doc_data[0].tensors)
+
+        # results ARE the same if we call from_disk
+        nlp2.from_disk(file_path)
+        doc2 = nlp2("We're interested at underwater basket weaving.")
+        tagger2 = nlp2.get_pipe("tagger")
+        tagger_trf2 = tagger2.model.get_ref("tok2vec").layers[0]
+        doc_tensor2 = tagger_trf2.predict([doc2])
+        assert_equal(doc_tensor2.doc_data[0].tensors, doc_tensor.doc_data[0].tensors)