from typing import List import pytest from thinc.api import fix_random_seed, Adam, set_dropout_rate from numpy.testing import assert_array_equal import numpy from spacy.ml.models import build_Tok2Vec_model from spacy.ml.models import build_text_classifier, build_simple_cnn_text_classifier from spacy.lang.en import English from spacy.lang.en.examples import sentences as EN_SENTENCES def get_all_params(model): params = [] for node in model.walk(): for name in node.param_names: params.append(node.get_param(name).ravel()) return node.ops.xp.concatenate(params) def get_docs(): nlp = English() return list(nlp.pipe(EN_SENTENCES + [" ".join(EN_SENTENCES)])) def get_gradient(model, Y): if isinstance(Y, model.ops.xp.ndarray): dY = model.ops.alloc(Y.shape, dtype=Y.dtype) dY += model.ops.xp.random.uniform(-1.0, 1.0, Y.shape) return dY elif isinstance(Y, List): return [get_gradient(model, y) for y in Y] else: raise ValueError(f"Could not get gradient for type {type(Y)}") def test_tok2vec(): return build_Tok2Vec_model(**TOK2VEC_KWARGS) TOK2VEC_KWARGS = { "width": 96, "embed_size": 2000, "subword_features": True, "char_embed": False, "conv_depth": 4, "bilstm_depth": 0, "maxout_pieces": 4, "window_size": 1, "dropout": 0.1, "nM": 0, "nC": 0, "pretrained_vectors": None, } TEXTCAT_KWARGS = { "width": 64, "embed_size": 2000, "pretrained_vectors": None, "exclusive_classes": False, "ngram_size": 1, "window_size": 1, "conv_depth": 2, "dropout": None, "nO": 7, } TEXTCAT_CNN_KWARGS = { "tok2vec": test_tok2vec(), "exclusive_classes": False, "nO": 13, } @pytest.mark.parametrize( "seed,model_func,kwargs", [ (0, build_Tok2Vec_model, TOK2VEC_KWARGS), (0, build_text_classifier, TEXTCAT_KWARGS), (0, build_simple_cnn_text_classifier, TEXTCAT_CNN_KWARGS), ], ) def test_models_initialize_consistently(seed, model_func, kwargs): fix_random_seed(seed) model1 = model_func(**kwargs) model1.initialize() fix_random_seed(seed) model2 = model_func(**kwargs) model2.initialize() params1 = get_all_params(model1) params2 = get_all_params(model2) assert_array_equal(params1, params2) @pytest.mark.parametrize( "seed,model_func,kwargs,get_X", [ (0, build_Tok2Vec_model, TOK2VEC_KWARGS, get_docs), (0, build_text_classifier, TEXTCAT_KWARGS, get_docs), (0, build_simple_cnn_text_classifier, TEXTCAT_CNN_KWARGS, get_docs), ], ) def test_models_predict_consistently(seed, model_func, kwargs, get_X): fix_random_seed(seed) model1 = model_func(**kwargs).initialize() Y1 = model1.predict(get_X()) fix_random_seed(seed) model2 = model_func(**kwargs).initialize() Y2 = model2.predict(get_X()) if model1.has_ref("tok2vec"): tok2vec1 = model1.get_ref("tok2vec").predict(get_X()) tok2vec2 = model2.get_ref("tok2vec").predict(get_X()) for i in range(len(tok2vec1)): for j in range(len(tok2vec1[i])): assert_array_equal( numpy.asarray(tok2vec1[i][j]), numpy.asarray(tok2vec2[i][j]) ) if isinstance(Y1, numpy.ndarray): assert_array_equal(Y1, Y2) elif isinstance(Y1, List): assert len(Y1) == len(Y2) for y1, y2 in zip(Y1, Y2): assert_array_equal(y1, y2) else: raise ValueError(f"Could not compare type {type(Y1)}") @pytest.mark.parametrize( "seed,dropout,model_func,kwargs,get_X", [ (0, 0.2, build_Tok2Vec_model, TOK2VEC_KWARGS, get_docs), (0, 0.2, build_text_classifier, TEXTCAT_KWARGS, get_docs), (0, 0.2, build_simple_cnn_text_classifier, TEXTCAT_CNN_KWARGS, get_docs), ], ) def test_models_update_consistently(seed, dropout, model_func, kwargs, get_X): def get_updated_model(): fix_random_seed(seed) optimizer = Adam(0.001) model = model_func(**kwargs).initialize() initial_params = get_all_params(model) set_dropout_rate(model, dropout) for _ in range(5): Y, get_dX = model.begin_update(get_X()) dY = get_gradient(model, Y) get_dX(dY) model.finish_update(optimizer) updated_params = get_all_params(model) with pytest.raises(AssertionError): assert_array_equal(initial_params, updated_params) return model model1 = get_updated_model() model2 = get_updated_model() assert_array_equal(get_all_params(model1), get_all_params(model2))