[BUG] When using model.fit(... validation_data=...) explicit calls to model.evaluate() fails after (in graph mode)

[gabrielspmoreira]

Bug description

The the TF model is compiled in graph mode, abd model.evaluate() is called by model.fit() when the validation_data is set for fit(), running validation at the end of each epoch.
But when model.evaluate() is explicitly called by the user after model.fit(validation_data=...) an error is raised.

Steps/Code to reproduce bug

1. If this PR is not merged yet, use its branch metrics_opt3 to reproduce the issue
2. Uncomment these lines on test_two_tower_retrieval_model_with_metrics and run it with run_eagerly=True
3. The explicitly call to model.evaluate() in the end of the test will fail

Expected behavior

It should compute both validation metrics during training and also when model.evaluate() is called

Additional context

I believe the issue happens because when we call model.fit() the ItemRetrievalTask._pre_eval_topk is not defined yet. The _pre_eval_topk is only set after model.fit(), when the model.model.load_topk_evaluation() is called (like here) to use the items dataset and the trained item tower.

The issue is that when model.fit(validation_data=...) is called first (and internally model.evaluation()), as _pre_eval_topk is frozen in the graph as None instead of being TopKIndexBlock, it is not included in the evaluation graph failing in the explicit model.evaluate().

Potential solution

A potential solution could be setting _pre_eval_topk=TopKIndexBlock() automatically for the ItemRetrievalTask, with dynamic variables ids and values with the first dim undefined, and than loading the data into the index after fit using something like variable.assign()

[sararb]

Thanks for the catch @gabrielspmoreira ! The main motivation of introducing load_topk_evaluation step is to allow the user for selecting the dataset of negative items to use for a Brute Force evaluation (and as further improvement we could support other TopKIndexBlock such as ANN-based).

So one possible solution, is to expose the set of evaluation candidates evaluation_candidates as an argument of the RetrievalModel and add a step to .evaluate() method that checks if pre_eval_topk is set. If not, it will load the index from the evaluation_candidates and set the pre_eval_topk accordingly.

The code would look like this:

    # Setting up evaluation
    item_features = ecommerce_data.schema.select_by_tag(Tags.ITEM).column_names
    item_dataset = ecommerce_data.dataframe[item_features].drop_duplicates()
    item_dataset = Dataset(item_dataset)
    model.evaluation_candidates = item_dataset
    model.evaluation_top_k = 10

    model.compile(optimizer="adam", run_eagerly=run_eagerly)

    # Training
    num_epochs = 2
    losses = model.fit(
        ecommerce_data.tf_dataloader(batch_size=50),
        epochs=num_epochs,
        train_metrics_steps=3,
        validation_data=ecommerce_data.tf_dataloader(batch_size=50),
        validation_steps=3,
    )