View source on GitHub
An ItemSelector that selects the first n items in the batch.
text.FirstNItemSelector(
num_to_select, unselectable_ids=None
)
Args | |
|---|---|
num_to_select
|
An int which is the leading number of items to select. |
unselectable_ids
|
(optional) A list of int ids that cannot be selected. Default is empty list. |
Methods
get_selectable
get_selectable(
input_ids, axis
)
See get_selectable() in superclass.
get_selection_mask
get_selection_mask(
input_ids, axis=1
)
Returns a mask of items that have been selected.
The default implementation simply returns all items not excluded by
get_selectable.
| Args | |
|---|---|
input_ids
|
A RaggedTensor.
|
axis
|
(optional) An int detailing the dimension to apply selection on. Default is the 1st dimension. |
| Returns | |
|---|---|
a RaggedTensor with shape input_ids.shape[:axis]. Its values are True
if the corresponding item (or broadcasted subitems) should be selected.
|