text.RoundRobinTrimmer

View source on GitHub

A Trimmer that allocates a length budget to segments via round robin.

Inherits From: Trimmer

text.RoundRobinTrimmer(
    max_seq_length, axis=-1
)

Used in the notebooks

Used in the guide
BERT Preprocessing with TF Text

A Trimmer that allocates a length budget to segments using a round robin strategy, then drops elements outside of the segment's allocated budget. See generate_mask() for more details.

Args
max_seq_length	a scalar Tensor int32 that describes the number max number of elements allowed in a batch.
axis	Axis to apply trimming on.

Methods

generate_mask

View source

generate_mask(
    segments
)

Calculates a truncation mask given a per-batch budget.

Calculate a truncation mask given a budget of the max number of items for each or all batch row. The allocation of the budget is done using a 'round robin' algorithm. This algorithm allocates quota in each bucket, left-to-right repeatedly until all the buckets are filled.

For example if the budget of [5] and we have segments of size [3, 4, 2], the truncate budget will be allocated as [2, 2, 1].

Args
segments	A list of RaggedTensor each w/ a shape of [num_batch, (num_items)].

Returns
a list with len(segments) of RaggedTensors, see superclass for details.

trim

View source

trim(
    segments
)

Truncate the list of segments.

Truncate the list of segments using the truncation strategy defined by generate_mask.

Args
segments	A list of RaggedTensors w/ shape [num_batch, (num_items)].

Returns
a list of RaggedTensors with len(segments) number of items and where each item has the same shape as its counterpart in segments and with unwanted values dropped. The values are dropped according to the TruncationStrategy defined.