Output Batch instead

Author: marcromeyn

merlin/models/torch/batch.py

@@ -14,7 +14,7 @@
 # limitations under the License.
 #
 
-from typing import Dict, Optional, Tuple, Union
+from typing import Dict, Optional, Union
 
 import torch
 
@@ -213,7 +213,7 @@ def sample_batch(
     dataset_or_loader: Union[Dataset, Loader],
     batch_size: Optional[int] = None,
     shuffle: Optional[bool] = False,
-) -> Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor]]:
+) -> Batch:
     """Util function to generate a batch of input tensors from a merlin.io.Dataset instance
 
     Parameters
@@ -244,7 +244,7 @@ def sample_batch(
     # batch could be of type Prediction, so we can't unpack directly
     inputs, targets = batch[0], batch[1]
 
-    return inputs, targets
+    return Batch(inputs, targets)
 
 
 def sample_features(
@@ -269,4 +269,4 @@ def sample_features(
         dictionary of feature tensors.
     """
 
-    return sample_batch(dataset_or_loader, batch_size, shuffle)[0]
+    return sample_batch(dataset_or_loader, batch_size, shuffle).features

tests/unit/torch/test_batch.py

@@ -162,31 +162,31 @@ class Test_sample_batch:
     def test_loader(self, music_streaming_data):
         loader = Loader(music_streaming_data, batch_size=2)
 
-        features, targets = sample_batch(loader)
+        batch = sample_batch(loader)
 
-        assert isinstance(features, dict)
-        assert len(list(features.keys())) == 12
-        for key, val in features.items():
+        assert isinstance(batch.features, dict)
+        assert len(list(batch.features.keys())) == 12
+        for key, val in batch.features.items():
             if not key.endswith("__values") and not key.endswith("__offsets"):
                 assert val.shape[0] == 2
 
-        assert isinstance(targets, dict)
-        assert list(targets.keys()) == ["click", "play_percentage", "like"]
-        for val in targets.values():
+        assert isinstance(batch.targets, dict)
+        assert list(batch.targets.keys()) == ["click", "play_percentage", "like"]
+        for val in batch.targets.values():
             assert val.shape[0] == 2
 
     def test_dataset(self, music_streaming_data):
-        features, targets = sample_batch(music_streaming_data, batch_size=2)
+        batch = sample_batch(music_streaming_data, batch_size=2)
 
-        assert isinstance(features, dict)
-        assert len(list(features.keys())) == 12
-        for key, val in features.items():
+        assert isinstance(batch.features, dict)
+        assert len(list(batch.features.keys())) == 12
+        for key, val in batch.features.items():
             if not key.endswith("__values") and not key.endswith("__offsets"):
                 assert val.shape[0] == 2
 
-        assert isinstance(targets, dict)
-        assert list(targets.keys()) == ["click", "play_percentage", "like"]
-        for val in targets.values():
+        assert isinstance(batch.targets, dict)
+        assert list(batch.targets.keys()) == ["click", "play_percentage", "like"]
+        for val in batch.targets.values():
             assert val.shape[0] == 2