Wrap Features in Class for compatibility with ItemRetrievalScorer

merlin/models/tf/core/prediction.py

@@ -20,8 +20,16 @@
 TensorLike = Union[tf.Tensor, tf.SparseTensor, tf.RaggedTensor]
 
 
+class Features:
+    def __init__(self, values: Dict[str, TensorLike]):
+        self.values = values
+
+    def __getitem__(self, key: str) -> TensorLike:
+        return self.values[key]
+
+
 class PredictionContext(NamedTuple):
-    features: Dict[str, TensorLike]
+    features: Features
     targets: Optional[Union[tf.Tensor, Dict[str, tf.Tensor]]] = None
     mask: tf.Tensor = (None,)
     training: bool = False

merlin/models/tf/models/base.py

@@ -13,7 +13,7 @@
 import merlin.io
 from merlin.models.tf.core.base import Block, ModelContext, PredictionOutput, is_input_block
 from merlin.models.tf.core.combinators import SequentialBlock
-from merlin.models.tf.core.prediction import Prediction, PredictionContext
+from merlin.models.tf.core.prediction import Features, Prediction, PredictionContext
 from merlin.models.tf.core.tabular import TabularBlock
 from merlin.models.tf.core.transformations import AsDenseFeatures, AsRaggedFeatures
 from merlin.models.tf.dataset import BatchedDataset
@@ -730,7 +730,9 @@ def call(self, inputs, targets=None, training=False, testing=False, output_conte
     def _create_context(
         self, inputs, targets=None, training=False, testing=False
     ) -> PredictionContext:
-        context = PredictionContext(inputs, targets=targets, training=training, testing=testing)
+        context = PredictionContext(
+            Features(inputs), targets=targets, training=training, testing=testing
+        )
 
         return context