tfrs.layers.factorized_top_k.TopK

Interface for top K layers.

Implementers must provide the following two methods:

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

Methods

call

View source

Query the index.

Args
queries Query features.
k Number of candidates to retrieve.

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

Raises
ValueError if index has not been called.

index

View source

Builds the retrieval index.

Args
candidates Matrix (or dataset) of candidate embeddings.
identifiers Optional tensor (or dataset) of candidate identifiers. If given these will be return to identify top candidates when performing searches. If not given, indices into the candidates tensor will be given instead.