diff --git a/spacy/ml/models/coref.py b/spacy/ml/models/coref.py index 24b5500a2..5042c10da 100644 --- a/spacy/ml/models/coref.py +++ b/spacy/ml/models/coref.py @@ -16,10 +16,10 @@ def build_wl_coref_model( tok2vec: Model[List[Doc], List[Floats2d]], embedding_size: int = 20, hidden_size: int = 1024, - n_hidden_layers: int = 1, # TODO rename to "depth"? + depth: int = 1, dropout: float = 0.3, # pairs to keep per mention after rough scoring - rough_k: int = 50, + rough_candidates: int = 50, # TODO is this not a training loop setting? a_scoring_batch_size: int = 512, # span predictor embeddings @@ -38,9 +38,9 @@ def build_wl_coref_model( dim, embedding_size, hidden_size, - n_hidden_layers, + depth, dropout, - rough_k, + rough_candidates, a_scoring_batch_size, ), convert_inputs=convert_coref_scorer_inputs, diff --git a/spacy/ml/models/span_predictor.py b/spacy/ml/models/span_predictor.py index b990b4019..ea445913b 100644 --- a/spacy/ml/models/span_predictor.py +++ b/spacy/ml/models/span_predictor.py @@ -16,6 +16,7 @@ def build_span_predictor( tok2vec: Model[List[Doc], List[Floats2d]], hidden_size: int = 1024, dist_emb_size: int = 64, + prefix: str = "coref_head_clusters" ): # TODO fix this try: @@ -30,7 +31,7 @@ def build_span_predictor( convert_inputs=convert_span_predictor_inputs, ) # TODO use proper parameter for prefix - head_info = build_get_head_metadata("coref_head_clusters") + head_info = build_get_head_metadata(prefix) model = (tok2vec & head_info) >> span_predictor return model diff --git a/spacy/pipeline/coref.py b/spacy/pipeline/coref.py index 5237788cc..a8813b7a3 100644 --- a/spacy/pipeline/coref.py +++ b/spacy/pipeline/coref.py @@ -33,9 +33,9 @@ default_config = """ @architectures = "spacy.Coref.v1" embedding_size = 20 hidden_size = 1024 -n_hidden_layers = 1 +depth = 1 dropout = 0.3 -rough_k = 50 +rough_candidates = 50 a_scoring_batch_size = 512 sp_embedding_size = 64 diff --git a/spacy/pipeline/span_predictor.py b/spacy/pipeline/span_predictor.py index 50c2e4ec6..d0561054d 100644 --- a/spacy/pipeline/span_predictor.py +++ b/spacy/pipeline/span_predictor.py @@ -26,6 +26,7 @@ default_span_predictor_config = """ @architectures = "spacy.SpanPredictor.v1" hidden_size = 1024 dist_emb_size = 64 +prefix = coref_head_clusters [model.tok2vec] @architectures = "spacy.Tok2Vec.v2"