Fix the output type of Padding op (#145)

test padding+embedding op

Author: sararb

merlin/dataloader/ops/padding.py

@@ -11,23 +11,20 @@
 
 
 class Padding(BaseOperator):
-    """Create an operator that will apply a embedding table to supplied indices.
-    This operator allows the user to supply an id lookup table if the indices supplied
-    via the id_lookup_table.
+    """Create an operator that will apply right padding to a given sequence.
+    This operator pads the sequence with a specified padding value up to a specified padding size.
+    If the sequence is longer than the padding size,
+    it is truncated to the first `padding size` elements.
 
     Parameters
     ----------
-    embeddings : np.ndarray
-        numpy ndarray representing embedding values
-    lookup_key : str, optional
-        the name of the column that will be used as indices, by default "id"
-    embedding_name : str, optional
-        name of new column of embeddings, added to output, by default "embeddings"
-    id_lookup_table : np.array, optional
-        numpy array of values that represent embedding indices, by default None
+    padding_size : int
+        The target size for the padded sequence
+    padding_value : Union[int, float]
+        The value to be used for padding the sequence, by default 0
     """
 
-    def __init__(self, padding_size: int, padding_value: Union[int, float]):
+    def __init__(self, padding_size: int, padding_value: Union[int, float] = 0):
         self.padding_size = padding_size
         self.padding_value = padding_value
 
@@ -76,7 +73,7 @@ def pad_put_zeros(column, padding_size, padding_val):
     # account for zero prepend
     array_lib = cupy if column.device == Device.GPU else np
     num_rows = len(column.offsets) - 1
-    zeros = array_lib.zeros((num_rows, padding_size)).flatten()
+    zeros = array_lib.zeros((num_rows, padding_size)).flatten() + padding_val
     row_lengths = column.offsets[1:] - column.offsets[:-1]
     row_ranges = []
     starts = array_lib.arange(num_rows) * padding_size
@@ -85,4 +82,5 @@ def pad_put_zeros(column, padding_size, padding_val):
         row_ranges.extend(array_lib.arange(int(starts[idx]), int(ends[idx])))
     array_lib.put(zeros, row_ranges, column.values)
     zeros = array_lib.reshape(zeros, (num_rows, padding_size))
+    zeros = zeros.astype(column.dtype.element_type.value)
     return zeros

tests/unit/dataloader/test_embeddings.py

@@ -22,6 +22,7 @@
 from merlin.core.dispatch import HAS_GPU
 from merlin.dataloader.loader_base import LoaderBase as Loader  # noqa
 from merlin.dataloader.ops.embeddings import EmbeddingOperator
+from merlin.dataloader.ops.padding import Padding
 from merlin.io import Dataset
 from merlin.schema import Tags
 from merlin.table import TensorColumn, TensorTable
@@ -245,3 +246,31 @@ def test_embedding_np_dl_with_lookup_ragged(
         assert (embeddings_offs == id_offsets).all()
         full_len += int(batch[0]["embeddings"].shape[0])
     assert full_len == offsets.shape[0] - 1
+
+
+def test_embedding_with_padding():
+    max_length = 10
+    batch_size = 3
+    id_embeddings = np.random.rand(1000, 16)
+    df = pd.DataFrame(
+        {
+            "id": [[0, 1, 2], [3, 4], [5, 6, 7, 8]],
+        }
+    )
+
+    dataset = Dataset(df)
+    transform = (
+        ["id"]
+        >> Padding(padding_size=max_length, padding_value=0)
+        >> EmbeddingOperator(id_embeddings, lookup_key="id", embedding_name="id_embedding")
+    )
+    data_loader = Loader(
+        dataset,
+        batch_size=batch_size,
+        transforms=transform,
+        shuffle=False,
+    )
+    x, _ = data_loader.peek()
+    assert x["id"].values.shape == (batch_size, max_length)
+    assert x["id_embedding"].values.shape == (batch_size, max_length, 16)
+    assert data_loader.output_schema.column_names == ["id", "id_embedding"]