tfrs.layers.factorized_top_k.TopK

Interface for top K layers.

tfrs.layers.factorized_top_k.TopK(
    k: int, **kwargs
) -> None

Implementers must provide the following two methods:

index: takes a tensor of candidate embeddings and creates the retrieval index.
call: takes a tensor of queries and returns top K candidates for those queries.

Methods

`call`

View source

@abc.abstractmethod
call(
    queries: Union[tf.Tensor, Dict[Text, tf.Tensor]], k: Optional[int] = None
) -> Tuple[tf.Tensor, tf.Tensor]

Query the index.

Args
`queries`	Query features. If `query_model` was provided in the constructor, these can be raw query features that will be processed by the query model before performing retrieval. If `query_model` was not provided, these should be pre-computed query embeddings.
`k`	The number of candidates to retrieve. If not supplied, defaults to the `k` value supplied in the constructor.

Returns
Tuple of (top candidate scores, top candidate identifiers).

Raises
ValueError if `index` has not been called.

`index`

View source

@abc.abstractmethod
index(
    candidates: tf.Tensor, identifiers: Optional[tf.Tensor] = None
) -> 'TopK'

Builds the retrieval index.

When called multiple times the existing index will be dropped and a new one created.

Args
`candidates`	Matrix of candidate embeddings.
`identifiers`	Optional tensor of candidate identifiers. If given, these will be used as identifiers of top candidates returned when performing searches. If not given, indices into the candidates tensor will be returned instead.

Returns
Self.

`index_from_dataset`

View source

index_from_dataset(
    candidates: tf.data.Dataset
) -> 'TopK'

Builds the retrieval index.

When called multiple times the existing index will be dropped and a new one created.

Args
`candidates`	Dataset of candidate embeddings or (candidate identifier, candidate embedding) pairs. If the dataset returns tuples, the identifiers will be used as identifiers of top candidates returned when performing searches. If not given, indices into the candidates dataset will be given instead.

Returns
Self.

Raises
ValueError if the dataset does not have the correct structure.

`is_exact`

View source

@abc.abstractmethod
is_exact() -> bool

Indicates whether the results returned by the layer are exact.

Some layers may return approximate scores: for example, the ScaNN layer may return approximate results.

Returns
True if the layer returns exact results, and False otherwise.

`query_with_exclusions`

View source

@tf.function
query_with_exclusions(
    queries: Union[tf.Tensor, Dict[Text, tf.Tensor]],
    exclusions: tf.Tensor,
    k: Optional[int] = None
) -> Tuple[tf.Tensor, tf.Tensor]

Query the index.

Args
`queries`	Query features. If `query_model` was provided in the constructor, these can be raw query features that will be processed by the query model before performing retrieval. If `query_model` was not provided, these should be pre-computed query embeddings.
`exclusions`	`[query_batch_size, num_to_exclude]` tensor of identifiers to be excluded from the top-k calculation. This is most commonly used to exclude previously seen candidates from retrieval. For example, if a user has already seen items with ids "42" and "43", you could set exclude to `[["42", "43"]]`.
`k`	The number of candidates to retrieve. Defaults to constructor `k` parameter if not supplied.

Returns
Tuple of (top candidate scores, top candidate identifiers).

Raises
ValueError if `index` has not been called. ValueError if `queries` is not a tensor (after being passed through the query model).