Makes RetrievalModelV2 support item tower with transforms (e.g. pre-trained embeddings)

gabrielspmoreira commented 1 year ago

Currently the methods based on batch_predict() support only a Dataset as input. That prevents performing batch_predict() using a Loader, which might contain transforms that are important to the model
A particular use case not supported because of that limitation is generating the item embeddings from the item tower of RetrievalModelV2, when among item features there are pre-trained embeddings loaded with EmbeddingOperator transform.
This PR, adds support to Loader with transforms to all batch_predict() based methods (e.g. to_top_k_encoder().

The batch_predict() method and associated methods were changed to accept Loader besides Dataset. When a Loader is passed, it cascades any transforms (like the EmbeddingOperator) to the Loader that is created inside the data_iterator_func.
It also checks if the model schema matches the Loader output schema (which is different from the dataset schema).

Added the test_two_tower_model_topk_evaluation_with_pretrained_emb to ensure that it is possible to have pre-trained embeddings in the item tower and convert the retrieval_model.to_top_k_encoder().