Tensorflow Recommenders add candidate at inference time

Hello everyone, I’m working on a recommender engine using tfrs and I’m wundering if there is any possibility to add candidates dataset to be queryed at inference time in the retrieval layer?

What currently I have seen is:

# Create a model that takes in raw query features, and
index = tfrs.layers.factorized_top_k.BruteForce(model.user_model)
# recommends movies out of the entire movies dataset.
  tf.data.Dataset.zip((movies.batch(100), movies.batch(100).map(model.movie_model)))

# Get recommendations.
_, titles = index(tf.constant(["42"]))
print(f"Recommendations for user 42: {titles[0, :3]}")

But I would like to have this to:

# Get recommendations.
candidates = new Dataset not same and different from used dataset in bruteforce
_, titles = index(tf.constant(["42"]), candidates)
print(f"Recommendations for user 42: {titles[0, :3]}")

I think @Wei_Wei can help here