Make pipe base class a bit less presumptuous
@@ -83,8 +83,12 @@ class Pipe(object):
"""
for docs in util.minibatch(stream, size=batch_size):
docs = list(docs)
- scores, tensors = self.predict(docs)
+ predictions = self.predict(docs)
+ if isinstance(predictions, tuple) and len(tuple) == 2:
+ scores, tensors = predictions
self.set_annotations(docs, scores, tensor=tensors)
+ else:
+ self.set_annotations(docs, predictions)
yield from docs
def predict(self, docs):
@@ -133,6 +137,7 @@ class Pipe(object):
If no model has been initialized yet, the model is added."""
if self.model is True:
self.model = self.Model(**self.cfg)
+ if hasattr(self, "vocab"):
link_vectors_to_models(self.vocab)
if sgd is None:
sgd = self.create_optimizer()
@@ -153,6 +158,7 @@ class Pipe(object):
serialize["cfg"] = lambda: srsly.json_dumps(self.cfg)
if self.model not in (True, False, None):
serialize["model"] = self.model.to_bytes
+ if hasattr(self, "vocab"):
serialize["vocab"] = self.vocab.to_bytes
exclude = util.get_serialization_exclude(serialize, exclude, kwargs)
return util.to_bytes(serialize, exclude)
@@ -173,6 +179,7 @@ class Pipe(object):
deserialize = OrderedDict()
deserialize["cfg"] = lambda b: self.cfg.update(srsly.json_loads(b))
+ if hasattr(self, "vocab"):
deserialize["vocab"] = lambda b: self.vocab.from_bytes(b)
deserialize["model"] = load_model
exclude = util.get_serialization_exclude(deserialize, exclude, kwargs)
