Repository: NVIDIA-Merlin/NVTabular

Issue:
[FEA] Include bucket range as category stat on HashBucket Op
DL model training workflows that are agnostic to the preprocessing used to generate a particular categorical feature will need to be able to recover the vocabulary size in order to construct the embedding. For features generated using the `Categorify` op, this is not an issue, since that information is recorded in the `Workflow` stats, which can be exported after preprocessing and read back in at train time.

But that information isn't recorded for the `HashBucket` op, which means you need to write some sort of exception which can somehow check for how that feature was generated then access the `Workflow` which was used to generate it, which is overall just a terrible way to have to implement things.

However, if we record `range(num_buckets)` as a stat for the op and mimic the API used for recovering the `Categorify` categories, then we can use the stats as an input to the training script which can then recover all the necessary information in an agnostic fashion.

Related to #78, which will combine these functionalities. Any implementation there should include the indices we include for the hash bins as part of the category stats.


Discussion:
Ah, you beat me to it :-)  @benfred let's merge these tickets.  I want to make sure that all forms of categorical encoding are covered, but Alec's description is way more detailed than mine.

Resolving patch:
diff --git a/nvtabular/ops/categorify.py b/nvtabular/ops/categorify.py
--- a/nvtabular/ops/categorify.py
+++ b/nvtabular/ops/categorify.py
@@ -378,6 +378,11 @@ def get_embedding_sizes(workflow):
     # when only hashing is applied, this will return embedding shape as (num_buckets, emb_dim)
     elif "buckets" in workflow.stats.keys():
         buckets = workflow.stats["buckets"]
+
+    # if we have hash buckets, but no coategorical just use the buckets
+    if buckets and "categories" not in workflow.stats:
+        return {col: _emb_sz_rule(num_rows) for col, num_rows in buckets.items()}
+
     if "mh" not in workflow.columns_ctx["categorical"]:
         return _get_embeddings_dask(workflow.stats["categories"], cols, buckets, freq)
     else:
diff --git a/nvtabular/ops/hash_bucket.py b/nvtabular/ops/hash_bucket.py
--- a/nvtabular/ops/hash_bucket.py
+++ b/nvtabular/ops/hash_bucket.py
@@ -17,12 +17,12 @@
 from cudf.utils.dtypes import is_list_dtype
 from nvtx import annotate
 
-from .categorify import _encode_list_column
+from .categorify import SetBuckets, _encode_list_column
 from .operator import CAT
-from .transform_operator import TransformOperator
+from .transform_operator import DFOperator
 
 
-class HashBucket(TransformOperator):
+class HashBucket(DFOperator):
     """
     This op maps categorical columns to a contiguous integer range
     by first hashing the column then modulating by the number of
@@ -127,3 +127,12 @@ def op_logic(self, gdf: cudf.DataFrame, target_columns: list, stats_context=None
 
             new_gdf[new_col] = encoded
         return new_gdf
+
+    @property
+    def req_stats(self):
+        return [
+            SetBuckets(
+                columns=self.columns,
+                num_buckets=self.num_buckets,
+            )
+        ]


Test patch:
diff --git a/tests/unit/test_ops.py b/tests/unit/test_ops.py
--- a/tests/unit/test_ops.py
+++ b/tests/unit/test_ops.py
@@ -376,6 +376,9 @@ def test_hash_bucket_lists(tmpdir):
     assert authors[0][0] == authors[1][0]  # 'User_A'
     assert authors[2][1] == authors[3][0]  # 'User_C'
 
+    # make sure we get the embedding sizes
+    assert nvt.ops.get_embedding_sizes(processor)["Authors"][0] == 10
+
 
 @pytest.mark.parametrize("engine", ["parquet"])
 def test_fill_missing(tmpdir, df, dataset, engine):