mirror of https://github.com/explosion/spaCy.git
52 lines
1.4 KiB
Python
52 lines
1.4 KiB
Python
from thinc.api import Model, noop
|
|
|
|
from ..util import registry
|
|
from .parser_model import ParserStepModel
|
|
|
|
|
|
@registry.layers("spacy.TransitionModel.v1")
|
|
def TransitionModel(
|
|
tok2vec, lower, upper, resize_output, dropout=0.2, unseen_classes=set()
|
|
):
|
|
"""Set up a stepwise transition-based model"""
|
|
if upper is None:
|
|
has_upper = False
|
|
upper = noop()
|
|
else:
|
|
has_upper = True
|
|
# don't define nO for this object, because we can't dynamically change it
|
|
return Model(
|
|
name="parser_model",
|
|
forward=forward,
|
|
dims={"nI": tok2vec.maybe_get_dim("nI")},
|
|
layers=[tok2vec, lower, upper],
|
|
refs={"tok2vec": tok2vec, "lower": lower, "upper": upper},
|
|
init=init,
|
|
attrs={
|
|
"has_upper": has_upper,
|
|
"unseen_classes": set(unseen_classes),
|
|
"resize_output": resize_output,
|
|
},
|
|
)
|
|
|
|
|
|
def forward(model, X, is_train):
|
|
step_model = ParserStepModel(
|
|
X,
|
|
model.layers,
|
|
unseen_classes=model.attrs["unseen_classes"],
|
|
train=is_train,
|
|
has_upper=model.attrs["has_upper"],
|
|
)
|
|
|
|
return step_model, step_model.finish_steps
|
|
|
|
|
|
def init(model, X=None, Y=None):
|
|
model.get_ref("tok2vec").initialize(X=X)
|
|
lower = model.get_ref("lower")
|
|
lower.initialize()
|
|
if model.attrs["has_upper"]:
|
|
statevecs = model.ops.alloc2f(2, lower.get_dim("nO"))
|
|
model.get_ref("upper").initialize(X=statevecs)
|