Fix per-sentence batching
from collections import namedtuple
from spacy.pipeline import Pipe
-from thinc.api import chain, foreach_sentence, layerize, wrap
+from thinc.api import chain, layerize, wrap
from thinc.neural.ops import get_array_module
from spacy.util import minibatch
@@ -57,12 +57,9 @@ class PyTT_TokenVectorEncoder(Pipe):
else:
model = PyTT_Wrapper(name)
nO = model.nO
- # TODO: Fix tensor if batch_by_length=0
batch_by_length = cfg.get("batch_by_length", 1)
- if batch_by_length:
model = with_length_batching(model, batch_by_length)
model = chain(get_word_pieces, model)
- if cfg.get("per_sentence", False):
model = foreach_sentence(model)
model.nO = nO
return model
@@ -228,3 +225,51 @@ def with_length_batching(model, min_batch):
return outputs, backprop_batched
return wrap(apply_model_to_batches, model)
+
+
+def foreach_sentence(layer, drop_factor=1.0):
+ """Map a layer across sentences (assumes spaCy-esque .sents interface)"""
+
+ def sentence_fwd(docs, drop=0.0):
+ sents = []
+ lengths = []
+ for doc in docs:
+ if doc.is_sentenced:
+ doc_sents = [sent for sent in doc.sents if len(sent)]
+ else:
+ doc_sents = [doc]
+ sents.extend(doc_sents)
+ lengths.append(len(doc_sents))
+ flat, bp_flat = layer.begin_update(sents, drop=drop)
+ outputs = _unflatten_ntuple_batch(layer.ops, flat, lengths)
+
+ def sentence_bwd(d_output, sgd=None):
+ d_flat = bp_flat(layer.ops.flatten(d_output), sgd=sgd)
+ if d_flat is None:
+ return d_sents
+ return _unflatten_ntuple_batch(layer.ops, d_flat, lengths)
+ return outputs, sentence_bwd
+
+ model = wrap(sentence_fwd, layer)
+ return model
+
+
+def _unflatten_ntuple_batch(ops, sent_outputs, lengths):
+ doc_outputs = []
+ offset = 0
+ for length in lengths:
+ doc_sents = sent_outputs[offset : offset + length]
+ # Transpose the nested list, from by-sentence to by-column.
+ doc_cols = []
+ for col_values in zip(*doc_sents):
+ # Merge the column values if possible
+ if len(col_values) and isinstance(col_values[0], ops.xp.ndarray):
+ doc_cols.append(ops.xp.vstack(col_values))
+ else:
+ doc_cols.append([])
+ for value in col_values:
+ doc_cols.extend(value)
+ # Now we can finally make the namedtuple
+ doc_outputs.append(sent_outputs[0]._make(doc_cols))
+ offset += length
+ return doc_outputs
