def ngram_suggester(docs: Iterable[Doc], *, ops: Optional[Ops] = None) -> Ragged:
if ops is None:
ops = get_current_ops()
spans = []
lengths = []
for doc in docs:
starts = ops.xp.arange(len(doc), dtype="i")
starts = starts.reshape((-1, 1))
length = 0
for size in sizes:
if size <= len(doc):
starts_size = starts[: len(doc) - (size - 1)]
spans.append(ops.xp.hstack((starts_size, starts_size + size)))
length += spans[-1].shape[0]
if spans:
assert spans[-1].ndim == 2, spans[-1].shape
lengths.append(length)
lengths_array = cast(Ints1d, ops.asarray(lengths, dtype="i"))
if len(spans) > 0:
output = Ragged(ops.xp.vstack(spans), lengths_array)
else:
output = Ragged(ops.xp.zeros((0, 0), dtype="i"), lengths_array)
assert output.dataXd.ndim == 2
return output
def forward(model: Model, source_spans: Tuple[Ragged, Ragged],
is_train: bool) -> Tuple[Ragged, Callable]:
"""Get subsequences from source vectors."""
ops = model.ops
X, spans = source_spans
assert spans.dataXd.ndim == 2
indices = _get_span_indices(ops, spans, X.lengths)
if len(indices) > 0:
Y = Ragged(X.dataXd[indices], spans.dataXd[:, 1] -
spans.dataXd[:, 0]) # type: ignore[arg-type, index]
else:
Y = Ragged(
ops.xp.zeros(X.dataXd.shape, dtype=X.dataXd.dtype),
ops.xp.zeros((len(X.lengths), ), dtype="i"),
)
x_shape = X.dataXd.shape
x_lengths = X.lengths
def backprop_windows(dY: Ragged) -> Tuple[Ragged, Ragged]:
dX = Ragged(ops.alloc2f(*x_shape), x_lengths)
ops.scatter_add(dX.dataXd, indices,
dY.dataXd) # type: ignore[arg-type]
return (dX, spans)
return Y, backprop_windows
def zeros(cls, length: int, width: int, *, xp=numpy) -> "TransformerData":
"""Create a valid TransformerData container for a given shape, filled
with zeros."""
return cls(tokens={"input_ids": numpy.zeros((length, ), dtype="i")},
tensors=[xp.zeros((1, length, width), dtype="f")],
align=Ragged(numpy.arange(length),
numpy.ones((length, ), dtype="i")))
def _truncate_alignment(align: Ragged, mask: numpy.ndarray) -> Ragged:
# We're going to have fewer wordpieces in the new array, so all of our
# wordpiece indices in the alignment table will be off --- they'll point
# to the wrong row. So we need to do three things here:
#
# 1) Adjust all the indices in align.dataXd to account for the dropped data
# 2) Remove the dropped indices from the align.dataXd
# 3) Calculate new align.lengths
#
# The wordpiece mapping is easily calculated by the cumulative sum of the
# mask table.
# Let's say we have [True, False, False, True]. The mapping of the dropped
# wordpieces doesn't matter, because we can filter it with the mask. So we
# have [0, 0, 0, 1], i.e the wordpiece that was
# at 0 is still at 0, and the wordpiece that was at 3 is now at 1.
mask = mask.ravel()
idx_map = mask.cumsum() - 1
idx_map[~mask] = -1
# Step 1: Adjust all the indices in align.dataXd.
new_align = idx_map[align.data.ravel()]
# Step 2: Remove the dropped indices
new_align = new_align[new_align >= 0]
# Step 3: Calculate new align.lengths
new_lengths = align.lengths.copy()
for i in range(len(align.lengths)):
drops = ~mask[align[i].data.ravel()]
new_lengths[i] -= drops.sum()
return Ragged(new_align, new_lengths)
def empty(cls) -> "TransformerData":
align = Ragged(numpy.zeros((0, ), dtype="i"),
numpy.zeros((0, ), dtype="i"))
return cls(wordpieces=WordpieceBatch.empty(),
tensors=[],
align=align,
attention=None)
def _trim_ragged_forward(model, Xr, is_train):
def backprop(dYr):
dY = dYr.data
dX = model.ops.alloc2f(dY.shape[0], dY.shape[1] + 1)
return Ragged(dX, dYr.lengths)
return Ragged(Xr.data[:, :-1], Xr.lengths), backprop
def zeros(cls, length: int, width: int, *, xp=numpy) -> "TransformerData":
"""Create a valid TransformerData container for a given shape, filled
with zeros."""
return cls(
wordpieces=WordpieceBatch.zeros([length], xp=xp),
tensors=[xp.zeros((1, length, width), dtype="f")],
align=Ragged(numpy.arange(length), numpy.ones((length, ),
dtype="i")),
)
def align(sequences):
lengths = []
indices = []
offset = 0
for seq in sequences:
for token_length in seq:
lengths.append(token_length)
indices.extend(i + offset for i in range(token_length))
offset += token_length
return Ragged(numpy.array(indices, dtype="i"), numpy.array(lengths, dtype="i"))
def ragged():
data = numpy.zeros((20, 4), dtype="f")
lengths = numpy.array([4, 2, 8, 1, 4], dtype="i")
data[0] = 0
data[1] = 1
data[2] = 2
data[3] = 3
data[4] = 4
data[5] = 5
return Ragged(data, lengths)
def test_ragged_empty():
data = numpy.zeros((0, 4), dtype="f")
lengths = numpy.array([], dtype="i")
ragged = Ragged(data, lengths)
assert_allclose(ragged[0:0].data, ragged.data)
assert_allclose(ragged[0:0].lengths, ragged.lengths)
assert_allclose(ragged[0:2].data, ragged.data)
assert_allclose(ragged[0:2].lengths, ragged.lengths)
assert_allclose(ragged[1:2].data, ragged.data)
assert_allclose(ragged[1:2].lengths, ragged.lengths)
def empty(cls, nr_docs) -> "FullTransformerBatch":
spans = [[] for i in range(nr_docs)]
doc_data = [TransformerData.empty() for i in range(nr_docs)]
align = Ragged(numpy.zeros((0, ), dtype="i"),
numpy.zeros((0, ), dtype="i"))
return cls(spans=spans,
tokens={},
tensors=[],
align=align,
_doc_data=doc_data)
def get_alignment(spans: List[Span], wordpieces: List[List[str]]) -> Ragged:
"""Compute a ragged alignment array that records, for each unique token in
`spans`, the corresponding indices in the flattened `wordpieces` array.
For instance, imagine you have two overlapping spans:
[[I, like, walking], [walking, outdoors]]
And their wordpieces are:
[[I, like, walk, ing], [walk, ing, out, doors]]
We want to align "walking" against [walk, ing, walk, ing], which have
indices [2, 3, 4, 5] once the nested wordpieces list is flattened.
The nested alignment list would be:
[[0], [1], [2, 3, 4, 5], [6, 7]]
I like walking outdoors
Which gets flattened into the ragged array:
[0, 1, 2, 3, 4, 5, 6, 7]
[1, 1, 4, 2]
The ragged format allows the aligned data to be computed via:
tokens = Ragged(wp_tensor[align.data], align.lengths)
This produces a ragged format, indicating which tokens need to be collapsed
to make the aligned array. The reduction is deferred for a later step, so
the user can configure it. The indexing is especially efficient in trivial
cases like this where the indexing array is completely continuous.
"""
if len(spans) != len(wordpieces):
raise ValueError("Cannot align batches of different sizes.")
# Tokens can occur more than once, and we need the alignment of each token
# to its place in the concatenated wordpieces array.
token_positions = get_token_positions(spans)
alignment: List[Set[int]] = [set() for _ in range(len(token_positions))]
wp_start = 0
for i, (span, wp_toks) in enumerate(zip(spans, wordpieces)):
sp_toks = [token.text for token in span]
span2wp, wp2span = tokenizations.get_alignments(sp_toks, wp_toks)
for token, wp_js in zip(span, span2wp):
position = token_positions[token]
alignment[position].update(wp_start + j for j in wp_js)
wp_start += len(wp_toks)
lengths: List[int] = []
flat: List[int] = []
for a in alignment:
lengths.append(len(a))
flat.extend(sorted(a))
align = Ragged(numpy.array(flat, dtype="i"), numpy.array(lengths,
dtype="i"))
return align
def _apply_empty_alignment(ops, align, X):
shape = X.shape
Y = Ragged(
ops.alloc2f(align.lengths.shape[0], X.shape[1]),
ops.alloc1i(align.lengths.shape[0]) + 1,
)
def backprop_null_alignment(dY: Ragged) -> Floats2d:
return ops.alloc2f(*shape)
return Y, backprop_null_alignment
def test_size_mismatch(Xs):
for reduce in [
reduce_first, reduce_last, reduce_max, reduce_mean, reduce_sum
]:
model = reduce()
lengths = model.ops.asarray([x.shape[0] for x in Xs], dtype="i")
X = Ragged(model.ops.flatten(Xs), lengths)
Y, backprop = model(X, is_train=True)
Y_bad = Y[:-1]
with pytest.raises(ValueError):
backprop(Y_bad)
def test_config_dataclasses():
@my_registry.cats("catsie.ragged")
def catsie_ragged(arg: Ragged):
return arg
data = numpy.zeros((20, 4), dtype="f")
lengths = numpy.array([4, 2, 8, 1, 4], dtype="i")
ragged = Ragged(data, lengths)
config = {"cfg": {"@cats": "catsie.ragged", "arg": ragged}}
result = my_registry.resolve(config)["cfg"]
assert isinstance(result, Ragged)
assert list(result._get_cumsums()) == [4, 6, 14, 15, 19]
def test_reduce_last(Xs):
model = reduce_last()
lengths = model.ops.asarray([x.shape[0] for x in Xs], dtype="i")
X = Ragged(model.ops.flatten(Xs), lengths)
Y, backprop = model(X, is_train=True)
assert isinstance(Y, numpy.ndarray)
assert Y.shape == (len(Xs), Xs[0].shape[1])
assert Y.dtype == Xs[0].dtype
assert list(Y[0]) == list(Xs[0][-1])
assert list(Y[1]) == list(Xs[1][-1])
dX = backprop(Y)
assert dX.dataXd.shape == X.dataXd.shape
def empty(cls, nr_docs) -> "FullTransformerBatch":
spans = [[] for i in range(nr_docs)]
doc_data = [TransformerData.empty() for i in range(nr_docs)]
align = Ragged(numpy.zeros((0, ), dtype="i"),
numpy.zeros((0, ), dtype="i"))
return cls(
spans=spans,
wordpieces=WordpieceBatch.empty(),
model_output=ModelOutput(),
align=align,
cached_doc_data=doc_data,
)
def empty(cls, nr_docs) -> "FullTransformerBatch":
spans = [[] for i in range(nr_docs)]
doc_data = [TransformerData.empty() for i in range(nr_docs)]
align = Ragged(numpy.zeros((0, ), dtype="i"),
numpy.zeros((0, ), dtype="i"))
return cls(
spans=spans,
wordpieces=WordpieceBatch.empty(),
tensors=[],
align=align,
attention=None,
cached_doc_data=doc_data,
)
def forward(model: Model[List[Doc], Ragged], docs: List[Doc],
is_train: bool) -> Tuple[Ragged, Callable]:
token_count = sum(len(doc) for doc in docs)
if not token_count:
return _handle_empty(model.ops, model.get_dim("nO"))
key_attr: int = model.attrs["key_attr"]
keys: Ints1d = model.ops.flatten(
cast(Sequence, [doc.to_array(key_attr) for doc in docs]))
vocab: Vocab = docs[0].vocab
W = cast(Floats2d, model.ops.as_contig(model.get_param("W")))
if vocab.vectors.mode == Mode.default:
V = cast(Floats2d, model.ops.asarray(vocab.vectors.data))
rows = vocab.vectors.find(keys=keys)
V = model.ops.as_contig(V[rows])
elif vocab.vectors.mode == Mode.floret:
V = cast(Floats2d, vocab.vectors.get_batch(keys))
V = model.ops.as_contig(V)
else:
raise RuntimeError(Errors.E896)
try:
vectors_data = model.ops.gemm(V, W, trans2=True)
except ValueError:
raise RuntimeError(Errors.E896)
if vocab.vectors.mode == Mode.default:
# Convert negative indices to 0-vectors
# TODO: more options for UNK tokens
vectors_data[rows < 0] = 0
output = Ragged(
vectors_data,
model.ops.asarray([len(doc) for doc in docs],
dtype="i") # type: ignore
)
mask = None
if is_train:
mask = _get_drop_mask(model.ops, W.shape[0],
model.attrs.get("dropout_rate"))
if mask is not None:
output.data *= mask
def backprop(d_output: Ragged) -> List[Doc]:
if mask is not None:
d_output.data *= mask
model.inc_grad(
"W",
model.ops.gemm(cast(Floats2d, d_output.data),
model.ops.as_contig(V),
trans1=True),
)
return []
return output, backprop
def test_noop_transforms(noop_models, ragged_input, padded_input, list_input):
# Make distinct backprop values,
# to check that the gradients get passed correctly
d_ragged = Ragged(ragged_input.data + 1, ragged_input.lengths)
d_padded = padded_input.copy()
d_padded.data += 1
d_list = [dx + 1 for dx in list_input]
for model in noop_models:
print(model.name)
check_transform_doesnt_change_noop_values(model, padded_input,
d_padded)
check_transform_doesnt_change_noop_values(model, list_input, d_list)
check_transform_doesnt_change_noop_values(model, ragged_input,
d_ragged)
def get_loss(
self, examples: Iterable[Example], spans_scores: Tuple[Ragged, Floats2d]
) -> Tuple[float, float]:
"""Find the loss and gradient of loss for the batch of documents and
their predicted scores.
examples (Iterable[Examples]): The batch of examples.
spans_scores: Scores representing the model's predictions.
RETURNS (Tuple[float, float]): The loss and the gradient.
DOCS: https://spacy.io/api/spancategorizer#get_loss
"""
spans, scores = spans_scores
spans = Ragged(
self.model.ops.to_numpy(spans.data), self.model.ops.to_numpy(spans.lengths)
)
label_map = {label: i for i, label in enumerate(self.labels)}
target = numpy.zeros(scores.shape, dtype=scores.dtype)
offset = 0
for i, eg in enumerate(examples):
# Map (start, end) offset of spans to the row in the d_scores array,
# so that we can adjust the gradient for predictions that were
# in the gold standard.
spans_index = {}
spans_i = spans[i].dataXd
for j in range(spans.lengths[i]):
start = int(spans_i[j, 0]) # type: ignore
end = int(spans_i[j, 1]) # type: ignore
spans_index[(start, end)] = offset + j
for gold_span in self._get_aligned_spans(eg):
key = (gold_span.start, gold_span.end)
if key in spans_index:
row = spans_index[key]
k = label_map[gold_span.label_]
target[row, k] = 1.0
# The target is a flat array for all docs. Track the position
# we're at within the flat array.
offset += spans.lengths[i]
target = self.model.ops.asarray(target, dtype="f") # type: ignore
# The target will have the values 0 (for untrue predictions) or 1
# (for true predictions).
# The scores should be in the range [0, 1].
# If the prediction is 0.9 and it's true, the gradient
# will be -0.1 (0.9 - 1.0).
# If the prediction is 0.9 and it's false, the gradient will be
# 0.9 (0.9 - 0.0)
d_scores = scores - target
loss = float((d_scores**2).sum())
return loss, d_scores
def test_reduce_mean(Xs):
Xs = [x * 1000 for x in Xs] # use large numbers for numeric stability
model = reduce_mean()
lengths = model.ops.asarray([x.shape[0] for x in Xs], dtype="i")
X = Ragged(model.ops.flatten(Xs), lengths)
Y, backprop = model(X, is_train=True)
assert isinstance(Y, numpy.ndarray)
assert Y.shape == (len(Xs), Xs[0].shape[1])
assert Y.dtype == Xs[0].dtype
assert numpy.all(Y[0] == Y[0][0]) # all values in row should be equal
assert Y[0][0] == Xs[0].mean()
assert numpy.all(Y[1] == Y[1][0])
assert Y[1][0] == Xs[1].mean()
dX = backprop(Y)
assert dX.dataXd.shape == X.dataXd.shape
def instance_forward(model: Model[List[Doc], Floats2d], docs: List[Doc],
is_train: bool) -> Tuple[Floats2d, Callable]:
pooling = model.get_ref("pooling")
tok2vec = model.get_ref("tok2vec")
get_instances = model.attrs["get_instances"]
all_instances = [get_instances(doc) for doc in docs]
tokvecs, bp_tokvecs = tok2vec(docs, is_train)
ents = []
lengths = []
for doc_nr, (instances, tokvec) in enumerate(zip(all_instances, tokvecs)):
token_indices = []
for instance in instances:
for ent in instance:
token_indices.extend([i for i in range(ent.start, ent.end)])
lengths.append(ent.end - ent.start)
ents.append(tokvec[token_indices])
lengths = cast(Ints1d, model.ops.asarray(lengths, dtype="int32"))
entities = Ragged(model.ops.flatten(ents), lengths)
pooled, bp_pooled = pooling(entities, is_train)
# Reshape so that pairs of rows are concatenated
relations = model.ops.reshape2f(pooled, -1, pooled.shape[1] * 2)
def backprop(d_relations: Floats2d) -> List[Doc]:
d_pooled = model.ops.reshape2f(d_relations, d_relations.shape[0] * 2,
-1)
d_ents = bp_pooled(d_pooled).data
d_tokvecs = []
ent_index = 0
for doc_nr, instances in enumerate(all_instances):
shape = tokvecs[doc_nr].shape
d_tokvec = model.ops.alloc2f(*shape)
count_occ = model.ops.alloc2f(*shape)
for instance in instances:
for ent in instance:
d_tokvec[ent.start:ent.end] += d_ents[ent_index]
count_occ[ent.start:ent.end] += 1
ent_index += ent.end - ent.start
d_tokvec /= count_occ + 0.00000000001
d_tokvecs.append(d_tokvec)
d_docs = bp_tokvecs(d_tokvecs)
return d_docs
return relations, backprop
def apply_alignment(ops: Ops, align: Ragged,
X: Floats2d) -> Tuple[Ragged, Callable]:
"""Align wordpiece data (X) to match tokens, and provide a callback to
reverse it.
This function returns a Ragged array, which represents the fact that one
token may be aligned against multiple wordpieces. It's a nested list,
concatenated with a lengths array to indicate the nested structure.
The alignment is also a Ragged array, where the lengths indicate how many
wordpieces each token is aligned against. The output ragged therefore has
the same lengths as the alignment ragged, which means the output data
also has the same number of data rows as the alignment. The size of the
lengths array indicates the number of tokens in the batch.
The actual alignment is a simple indexing operation:
for i, index in enumerate(align.data):
Y[i] = X[index]
Which is vectorized via numpy advanced indexing:
Y = X[align.data]
The inverse operation, for the backward pass, uses the 'scatter_add' op
because one wordpiece may be aligned against multiple tokens. So we need:
for i, index in enumerate(align.data):
X[index] += Y[i]
The addition wouldn't occur if we simply did `X[index] = Y`, so we use
the scatter_add op.
"""
if not align.lengths.sum():
return _apply_empty_alignment(ops, align, X)
shape = X.shape
indices = cast(Ints1d, align.dataXd)
Y = Ragged(X[indices], cast(Ints1d, ops.asarray(align.lengths)))
def backprop_apply_alignment(dY: Ragged) -> Floats2d:
assert dY.data.shape[0] == indices.shape[0]
dX = ops.alloc2f(*shape)
ops.scatter_add(dX, indices, cast(Floats2d, dY.dataXd))
return dX
return Y, backprop_apply_alignment
def forward(model: Model, source_spans: Tuple[Ragged, Ragged],
is_train: bool) -> Tuple[Ragged, Callable]:
"""Get subsequences from source vectors."""
ops = model.ops
X, spans = source_spans
assert spans.dataXd.ndim == 2
indices = _get_span_indices(ops, spans, X.lengths)
Y = Ragged(X.dataXd[indices], spans.dataXd[:, 1] - spans.dataXd[:, 0])
x_shape = X.dataXd.shape
x_lengths = X.lengths
def backprop_windows(dY: Ragged) -> Tuple[Ragged, Ragged]:
dX = Ragged(ops.alloc2f(*x_shape), x_lengths)
ops.scatter_add(dX.dataXd, indices, dY.dataXd)
return (dX, spans)
return Y, backprop_windows
def get_alignment_via_offset_mapping(spans: List[Span], token_data) -> Ragged:
# Tokens can occur more than once, and we need the alignment of each token
# to its place in the concatenated wordpieces array.
token_positions = get_token_positions(spans)
alignment: List[Set[int]] = [set() for _ in range(len(token_positions))]
wp_start = 0
for i, span in enumerate(spans):
for j, token in enumerate(span):
position = token_positions[token]
for char_idx in range(token.idx, token.idx + len(token)):
wp_j = token_data.char_to_token(i, char_idx)
if wp_j is not None:
alignment[position].add(wp_start + wp_j)
wp_start += len(token_data.input_ids[i])
lengths: List[int] = []
flat: List[int] = []
for a in alignment:
lengths.append(len(a))
flat.extend(sorted(a))
align = Ragged(numpy.array(flat, dtype="i"), numpy.array(lengths,
dtype="i"))
return align
def forward(model: Model[List[Doc], Ragged], docs: List[Doc],
is_train: bool) -> Tuple[Ragged, Callable]:
if not sum(len(doc) for doc in docs):
return _handle_empty(model.ops, model.get_dim("nO"))
key_attr = model.attrs["key_attr"]
W = cast(Floats2d, model.ops.as_contig(model.get_param("W")))
V = cast(Floats2d, model.ops.asarray(docs[0].vocab.vectors.data))
rows = model.ops.flatten(
[doc.vocab.vectors.find(keys=doc.to_array(key_attr)) for doc in docs])
try:
vectors_data = model.ops.gemm(model.ops.as_contig(V[rows]),
W,
trans2=True)
except ValueError:
raise RuntimeError(Errors.E896)
# Convert negative indices to 0-vectors (TODO: more options for UNK tokens)
vectors_data[rows < 0] = 0
output = Ragged(vectors_data,
model.ops.asarray([len(doc) for doc in docs], dtype="i"))
mask = None
if is_train:
mask = _get_drop_mask(model.ops, W.shape[0],
model.attrs.get("dropout_rate"))
if mask is not None:
output.data *= mask
def backprop(d_output: Ragged) -> List[Doc]:
if mask is not None:
d_output.data *= mask
model.inc_grad(
"W",
model.ops.gemm(d_output.data,
model.ops.as_contig(V[rows]),
trans1=True),
)
return []
return output, backprop
def ragged_input(ops, list_input):
lengths = numpy.array([len(x) for x in list_input], dtype="i")
if not list_input:
return Ragged(ops.alloc2f(0, 0), lengths)
else:
return Ragged(ops.flatten(list_input), lengths)
def backprop(dYr):
dY = dYr.data
dX = model.ops.alloc2f(dY.shape[0], dY.shape[1] + 1)
return Ragged(dX, dYr.lengths)
def _handle_empty(ops: Ops, nO: int):
return Ragged(ops.alloc2f(0, nO), ops.alloc1i(0)), lambda d_ragged: []
