Skip to content

Commit

Permalink
Support ragged tensor in seqio.evaluation in calculating the max seq …
Browse files Browse the repository at this point in the history
…length.

PiperOrigin-RevId: 674353863
  • Loading branch information
SeqIO Team authored and SeqIO committed Sep 17, 2024
1 parent 45013cc commit 958249d
Showing 1 changed file with 6 additions and 1 deletion.
7 changes: 6 additions & 1 deletion seqio/evaluation.py
Original file line number Diff line number Diff line change
Expand Up @@ -131,7 +131,12 @@ def _cache_and_measure_examples(
for ex in tfds.as_numpy(ds):
for k in max_sequence_length:
sequence_dim = sequence_dims.get(k, 0)
sequence_length = ex[k].shape[sequence_dim]
if isinstance(ex[k], tf.RaggedTensor):
sequence_length = tf.reduce_max(
ex[k].row_lengths(axis=sequence_dim)
).numpy()
else:
sequence_length = ex[k].shape[sequence_dim]
max_sequence_length[k] = max(max_sequence_length[k], sequence_length)
cnt += 1

Expand Down

0 comments on commit 958249d

Please sign in to comment.