tfrs.layers.factorized_top_k.Streaming

Retrieves K highest scoring items and their ids from a large dataset.

Inherits From: TopK

tfrs.layers.factorized_top_k.Streaming(
    k: int = 10,
    handle_incomplete_batches: bool = True,
    num_parallel_calls: int = tf.data.experimental.AUTOTUNE,
    sorted_order: bool = True
) -> None

Used to efficiently retrieve top K query-candidate scores from a dataset, along with the top scoring candidates' identifiers.

k Number of top scores to retrieve.
handle_incomplete_batches When True, candidate batches smaller than k will be correctly handled at the price of some performance. As an alternative, consider using the drop_remainer option when batching the candidate dataset.
num_parallel_calls Degree of parallelism when computing scores. Defaults to autotuning.
sorted_order If the resulting scores should be returned in sorted order. setting this to False may result in a small increase in performance.

ValueError if candidate elements are not tuples.

Methods

call

View source

call(
    query_embeddings: tf.Tensor,
    k: Optional[int] = None
) -> Tuple[tf.Tensor, tf.Tensor]

Computes K highest scores and candidate indices for a given query.

Args
query_embeddings [query_batch_size, embedding_dim] tensor of query embeddings.
k Number of elements to retrieve. If not set, will default to the k set in the constructor.

Returns
Tuple of [query_batch_size, k] tensor of top scores for each query and [query_batch_size, k] tensor of indices for highest scoring candidates.

index

View source

index(
    candidates: tf.data.Dataset,
    identifiers: Optional[tf.data.Dataset] = None
) -> "Streaming"

Sets the dataset of candidates over which to compute streaming top K.

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 datset will be given instead.

Returns
Self for chaining.