Can i export retrieval model with query and candidate as inputs?

General Discussion
,
#1

I have a factorized topK retrieval model, everything works fine, but I want to input specifics slices of candidates data, its possible? In the code i want to turn parsed_topk dataset in to a input of the model

index = tfrs.layers.factorized_top_k.BruteForce(final_model.query_model)

index.index_from_dataset(
    tf.data.Dataset.zip((parsed_topK.batch(128).map(lambda x: x['experience.jobTitle']), parsed_topK.batch(128).map(final_model.candidate_model)))
)