fixing arguments

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,

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

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
 

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"