diff --git a/spacy/ml/models/textcat.py b/spacy/ml/models/textcat.py index 3e5471ab3..601c94a7f 100644 --- a/spacy/ml/models/textcat.py +++ b/spacy/ml/models/textcat.py @@ -185,6 +185,11 @@ def build_text_classifier_v2( def init_ensemble_textcat(model, X, Y) -> Model: + # When tok2vec is lazily initialized, we need to initialize it before + # the rest of the chain to ensure that we can get its width. + tok2vec = model.get_ref("tok2vec") + tok2vec.initialize(X) + tok2vec_width = get_tok2vec_width(model) model.get_ref("attention_layer").set_dim("nO", tok2vec_width) model.get_ref("maxout_layer").set_dim("nO", tok2vec_width) @@ -264,6 +269,7 @@ def _build_parametric_attention_with_residual_nonlinear( parametric_attention.set_ref("tok2vec", tok2vec) parametric_attention.set_ref("attention_layer", attention_layer) + parametric_attention.set_ref("key_transform", key_transform) parametric_attention.set_ref("nonlinear_layer", nonlinear_layer) parametric_attention.set_ref("norm_layer", norm_layer) @@ -271,10 +277,17 @@ def _build_parametric_attention_with_residual_nonlinear( def _init_parametric_attention_with_residual_nonlinear(model, X, Y) -> Model: + # When tok2vec is lazily initialized, we need to initialize it before + # the rest of the chain to ensure that we can get its width. + tok2vec = model.get_ref("tok2vec") + tok2vec.initialize(X) + tok2vec_width = get_tok2vec_width(model) model.get_ref("attention_layer").set_dim("nO", tok2vec_width) - model.get_ref("nonlinear_layer").set_dim("nO", tok2vec_width) + model.get_ref("key_transform").set_dim("nI", tok2vec_width) + model.get_ref("key_transform").set_dim("nO", tok2vec_width) model.get_ref("nonlinear_layer").set_dim("nI", tok2vec_width) + model.get_ref("nonlinear_layer").set_dim("nO", tok2vec_width) model.get_ref("norm_layer").set_dim("nI", tok2vec_width) model.get_ref("norm_layer").set_dim("nO", tok2vec_width) init_chain(model, X, Y) diff --git a/spacy/tests/pipeline/test_textcat.py b/spacy/tests/pipeline/test_textcat.py index 7a78c3dac..8a0c1a976 100644 --- a/spacy/tests/pipeline/test_textcat.py +++ b/spacy/tests/pipeline/test_textcat.py @@ -28,6 +28,8 @@ from spacy.tokens import Doc, DocBin from spacy.training import Example from spacy.training.initialize import init_nlp +# Ensure that the architecture gets added to the registry. +from ..tok2vec import build_lazy_init_tok2vec as _ from ..util import make_tempdir TRAIN_DATA_SINGLE_LABEL = [ @@ -40,6 +42,13 @@ TRAIN_DATA_MULTI_LABEL = [ ("I'm confused but happy", {"cats": {"ANGRY": 0.0, "CONFUSED": 1.0, "HAPPY": 1.0}}), ] +lazy_init_model_config = """ +[model] +@architectures = "test.LazyInitTok2Vec.v1" +width = 96 +""" +LAZY_INIT_TOK2VEC_MODEL = Config().from_str(lazy_init_model_config)["model"] + def make_get_examples_single_label(nlp): train_examples = [] @@ -546,6 +555,34 @@ def test_error_with_multi_labels(): nlp.initialize(get_examples=lambda: train_examples) +# fmt: off +@pytest.mark.parametrize( + "name,textcat_config", + [ + # ENSEMBLE V2 + ("textcat_multilabel", {"@architectures": "spacy.TextCatEnsemble.v2", "tok2vec": LAZY_INIT_TOK2VEC_MODEL, "linear_model": {"@architectures": "spacy.TextCatBOW.v3", "exclusive_classes": False, "ngram_size": 1, "no_output_layer": False}}), + ("textcat", {"@architectures": "spacy.TextCatEnsemble.v2", "tok2vec": LAZY_INIT_TOK2VEC_MODEL, "linear_model": {"@architectures": "spacy.TextCatBOW.v3", "exclusive_classes": True, "ngram_size": 5, "no_output_layer": False}}), + # PARAMETRIC ATTENTION V1 + ("textcat", {"@architectures": "spacy.TextCatParametricAttention.v1", "tok2vec": LAZY_INIT_TOK2VEC_MODEL, "exclusive_classes": True}), + ("textcat_multilabel", {"@architectures": "spacy.TextCatParametricAttention.v1", "tok2vec": LAZY_INIT_TOK2VEC_MODEL, "exclusive_classes": False}), + # REDUCE + ("textcat", {"@architectures": "spacy.TextCatReduce.v1", "tok2vec": LAZY_INIT_TOK2VEC_MODEL, "exclusive_classes": True, "use_reduce_first": True, "use_reduce_last": True, "use_reduce_max": True, "use_reduce_mean": True}), + ("textcat_multilabel", {"@architectures": "spacy.TextCatReduce.v1", "tok2vec": LAZY_INIT_TOK2VEC_MODEL, "exclusive_classes": False, "use_reduce_first": True, "use_reduce_last": True, "use_reduce_max": True, "use_reduce_mean": True}), + ], +) +# fmt: on +def test_tok2vec_lazy_init(name, textcat_config): + # Check that we can properly initialize and use a textcat model using + # a lazily-initialized tok2vec. + nlp = English() + pipe_config = {"model": textcat_config} + textcat = nlp.add_pipe(name, config=pipe_config) + textcat.add_label("POSITIVE") + textcat.add_label("NEGATIVE") + nlp.initialize() + nlp.pipe(["This is a test."]) + + @pytest.mark.parametrize( "name,get_examples, train_data", [ diff --git a/spacy/tests/tok2vec.py b/spacy/tests/tok2vec.py new file mode 100644 index 000000000..7e7b689eb --- /dev/null +++ b/spacy/tests/tok2vec.py @@ -0,0 +1,36 @@ +from typing import List + +from thinc.api import Model +from thinc.types import Floats2d + +from spacy.tokens import Doc +from spacy.util import registry + + +@registry.architectures("test.LazyInitTok2Vec.v1") +def build_lazy_init_tok2vec(*, width: int) -> Model[List[Doc], List[Floats2d]]: + """tok2vec model of which the output size is only known after + initialization. This implementation does not output meaningful + embeddings, it is strictly for testing.""" + return Model( + "lazy_init_tok2vec", + lazy_init_tok2vec_forward, + init=lazy_init_tok2vec_init, + dims={"nO": None}, + attrs={"width": width}, + ) + + +def lazy_init_tok2vec_init(model: Model, X=None, Y=None): + width = model.attrs["width"] + model.set_dim("nO", width) + + +def lazy_init_tok2vec_forward(model: Model, X: List[Doc], is_train: bool): + width = model.get_dim("nO") + Y = [model.ops.alloc2f(len(doc), width) for doc in X] + + def backprop(dY): + return [] + + return Y, backprop