Add support of transformer-based retrieval models

[sararb]

Goals ⚽

• Support ragged queries and targets in the BruteForce(TopKLayer) class. This is needed to export a sequential session encoder.
• Support exporting the candidate embeddings from CategoricalOutput class.
• Make the bias term optional in the weight-tying layer EmbeddingTablePrediction. The default is false, as we wouldn't have access to the bias term if we export the query encoder to an ANN system for inference.

Testing Details 🔍

• Add a check of ragged_query in the unit testtest_brute_force_layer.
• Add use_bias option to the unit test test_last_item_prediction.

[github-actions]

Documentation preview

https://nvidia-merlin.github.io/models/review/pr-1128

[sararb]

This assumes that we evaluate only on the last item in the session (which is the default mode during inference too). We might need to extend it to evaluate other items in the sequence in future work.

[gabrielspmoreira]

Thanks @sararb. I was trying to understand why we would have 1 as the 2nd dim of inputs, but from your explanation I now understand the reason is that we have predictions only for the last position.
I think it would be useful to add this remark as a comment.

[rnyak]

@sararb I can save the xlnet model but I cannot load it back. I am getting an error. ValueError: The last dimension of the input shape of a Dense layer should be defined. Found None. Received: input_shape=(None, None)

do you think you can add a test for that in the unit test? also a test to showcase how one can do offline prediction? thanks.

[gabrielspmoreira]

Thanks @sararb. I was trying to understand why we would have 1 as the 2nd dim of inputs, but from your explanation I now understand the reason is that we have predictions only for the last position.
I think it would be useful to add this remark as a comment.

[rnyak]

Note: with making the bias term optional in the weight-tying layer EmbeddingTablePrediction, we make sure that the training and inference models have same score calculation, bcs contrastive output head is not exported during inference.