[review] SID: address tiankongdeguiji review on PR alibaba#538

- Metrics: switch the shared reconstruction metric to
  torchmetrics.MeanSquaredError (correct (preds, target) aggregation vs the
  biased mean-of-batch-means a MeanMetric gave), and add a UniqueRatio metric
  class (tzrec/metrics/unique_ratio.py) for codebook coverage instead of a
  MeanMetric fed a hand-computed ratio. BaseSidModel._update_unique_sid_ratio
  now delegates to it; the empty-batch guard + DDP reduction live in the
  metric. The RQ-VAE train-path mse moves to MeanSquaredError too.
- __init__.py: revert tzrec/modules/sid_generation/__init__.py to a bare
  package marker (no re-exports) and import ResidualKMeansQuantizer from its
  submodule in sid_rqkmeans, per "avoid adding to __init__.py".

Tests: tzrec/metrics/unique_ratio_test.py (ratio, empty no-op, mean over
batches); sid_model_test verifies the delegation.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: WhiteSwan1

tzrec/metrics/unique_ratio.py

@@ -0,0 +1,48 @@
+# Copyright (c) 2026, Alibaba Group;
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#    http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+from torchmetrics import Metric
+
+
+class UniqueRatio(Metric):
+    """Codebook-coverage metric: mean of per-batch (unique rows / batch size).
+
+    Each ``update`` counts the unique rows of a ``(B, n_layers)`` semantic-ID
+    code tensor and accumulates the per-batch ratio; ``compute`` returns the
+    running mean. Empty batches (``B == 0``, e.g. an empty final DDP/TorchRec
+    shard) are skipped. States reduce by ``sum`` across ranks.
+    """
+
+    higher_is_better = True
+    is_differentiable = False
+
+    def __init__(self, **kwargs) -> None:
+        super().__init__(**kwargs)
+        self.add_state("ratio_sum", default=torch.tensor(0.0), dist_reduce_fx="sum")
+        self.add_state("count", default=torch.tensor(0.0), dist_reduce_fx="sum")
+
+    def update(self, codes: torch.Tensor) -> None:
+        """Accumulate the unique-ratio of one batch of codes.
+
+        Args:
+            codes (Tensor): semantic-ID codes, shape (B, n_layers).
+        """
+        batch_size = codes.shape[0]
+        if batch_size == 0:
+            return
+        unique = torch.unique(codes, dim=0).shape[0]
+        self.ratio_sum += unique / batch_size
+        self.count += 1
+
+    def compute(self) -> torch.Tensor:
+        """Mean per-batch unique ratio (NaN before any non-empty update)."""
+        return self.ratio_sum / self.count

tzrec/metrics/unique_ratio_test.py

@@ -0,0 +1,40 @@
+# Copyright (c) 2026, Alibaba Group;
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#    http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import torch
+
+from tzrec.metrics.unique_ratio import UniqueRatio
+
+
+class UniqueRatioTest(unittest.TestCase):
+    def test_known_duplicates(self) -> None:
+        metric = UniqueRatio()
+        # 3 unique rows out of 4 -> 0.75.
+        metric.update(torch.tensor([[1, 2], [1, 2], [3, 4], [5, 6]]))
+        self.assertAlmostEqual(metric.compute().item(), 0.75, places=6)
+
+    def test_empty_batch_skipped(self) -> None:
+        metric = UniqueRatio()
+        metric.update(torch.empty(0, 3, dtype=torch.long))
+        self.assertEqual(metric.count.item(), 0.0)
+        self.assertTrue(torch.isnan(metric.compute()))
+
+    def test_mean_over_batches(self) -> None:
+        metric = UniqueRatio()
+        metric.update(torch.tensor([[1, 1], [1, 1]]))  # 1/2 = 0.5
+        metric.update(torch.tensor([[1, 1], [2, 2]]))  # 2/2 = 1.0
+        self.assertAlmostEqual(metric.compute().item(), 0.75, places=6)  # mean
+
+
+if __name__ == "__main__":
+    unittest.main()

tzrec/models/sid_model.py

@@ -18,6 +18,7 @@
 
 from tzrec.datasets.utils import BASE_DATA_GROUP, Batch
 from tzrec.features.feature import BaseFeature
+from tzrec.metrics.unique_ratio import UniqueRatio
 from tzrec.models.model import BaseModel
 from tzrec.protos.model_pb2 import ModelConfig
 
@@ -100,8 +101,8 @@ def init_metric(self) -> None:
         ``unique_sid_ratio``: codebook coverage = unique SIDs / batch size.
         Subclasses call ``super().init_metric()`` then add their extras.
         """
-        self._metric_modules["mse"] = torchmetrics.MeanMetric()
-        self._metric_modules["unique_sid_ratio"] = torchmetrics.MeanMetric()
+        self._metric_modules["mse"] = torchmetrics.MeanSquaredError()
+        self._metric_modules["unique_sid_ratio"] = UniqueRatio()
 
     def update_train_metric(
         self,
@@ -116,13 +117,12 @@ def update_train_metric(
         return
 
     def _update_unique_sid_ratio(self, codes: torch.Tensor) -> None:
-        """Update the codebook-coverage metric (unique SIDs / batch size).
+        """Update the codebook-coverage metric.
+
+        The unique-ratio math (and the empty-batch guard) lives in
+        :class:`~tzrec.metrics.unique_ratio.UniqueRatio`.
 
         Args:
             codes (Tensor): semantic-ID codes, shape (B, n_layers).
         """
-        B = codes.shape[0]
-        if B == 0:  # empty final shard under DDP/TorchRec
-            return
-        unique_sids = torch.unique(codes, dim=0).shape[0]
-        self._metric_modules["unique_sid_ratio"].update(unique_sids / B)
+        self._metric_modules["unique_sid_ratio"].update(codes)

tzrec/models/sid_model_test.py

@@ -12,32 +12,30 @@
 import unittest
 
 import torch
-import torchmetrics
 
+from tzrec.metrics.unique_ratio import UniqueRatio
 from tzrec.models.sid_model import BaseSidModel
 
 
 class UpdateUniqueSidRatioTest(unittest.TestCase):
-    """Unit-test the codebook-coverage metric helper.
+    """Verify ``_update_unique_sid_ratio`` delegates to the UniqueRatio metric.
 
-    ``_update_unique_sid_ratio`` is pure tensor logic over
-    ``self._metric_modules`` (no proto dependency), so it is testable without
-    a full ``BaseSidModel`` config — full model coverage waits on the concrete
-    subclasses in the follow-up PRs.
+    The unique-ratio math itself is covered in
+    ``tzrec/metrics/unique_ratio_test.py``; here we only check the wiring (no
+    proto dependency, so it is testable without a full ``BaseSidModel`` config).
     """
 
     def _bare_model(self) -> BaseSidModel:
         # Bypass __init__ (which needs a pipeline config); only the metric
         # module the helper touches needs to exist.
         model = BaseSidModel.__new__(BaseSidModel)
-        model._metric_modules = {"unique_sid_ratio": torchmetrics.MeanMetric()}
+        model._metric_modules = {"unique_sid_ratio": UniqueRatio()}
         return model
 
     def test_empty_batch_is_noop(self) -> None:
         model = self._bare_model()
         model._update_unique_sid_ratio(torch.empty(0, 3, dtype=torch.long))
-        # The B == 0 guard returns early -> no sample recorded.
-        self.assertEqual(model._metric_modules["unique_sid_ratio"].weight.item(), 0.0)
+        self.assertEqual(model._metric_modules["unique_sid_ratio"].count.item(), 0.0)
 
     def test_ratio_on_known_duplicates(self) -> None:
         model = self._bare_model()

tzrec/models/sid_rqkmeans.py

@@ -27,8 +27,10 @@
 from tzrec.datasets.utils import Batch
 from tzrec.features.feature import BaseFeature
 from tzrec.models.sid_model import BaseSidModel
-from tzrec.modules.sid_generation import ResidualKMeansQuantizer
 from tzrec.modules.sid_generation.kmeans import recon_diagnostics
+from tzrec.modules.sid_generation.residual_kmeans_quantizer import (
+    ResidualKMeansQuantizer,
+)
 from tzrec.protos.model_pb2 import ModelConfig
 from tzrec.utils import config_util
 from tzrec.utils.logging_util import logger
@@ -255,11 +257,15 @@ def update_metric(
             losses (dict, optional): a dict of loss.
         """
         if "input_embedding" in predictions:
-            mse, rel = recon_diagnostics(
+            _, rel = recon_diagnostics(
                 predictions["input_embedding"],
                 predictions["quantized"],
             )
-            self._metric_modules["mse"].update(mse)
+            # MeanSquaredError aggregates (preds, target) itself; rel_loss has
+            # no torchmetrics equivalent so it stays a MeanMetric.
+            self._metric_modules["mse"].update(
+                predictions["quantized"], predictions["input_embedding"]
+            )
             self._metric_modules["rel_loss"].update(rel)
 
         self._update_unique_sid_ratio(predictions["codes"])

tzrec/models/sid_rqvae.py

@@ -327,7 +327,7 @@ def init_metric(self) -> None:
         # is intentionally eval-only: torch.unique(codes, dim=0).shape[0]
         # forces a GPU->host sync every step, and codebook coverage is a
         # diagnostic, not a training signal.
-        self._train_metric_modules["mse"] = torchmetrics.MeanMetric()
+        self._train_metric_modules["mse"] = torchmetrics.MeanSquaredError()
 
     def update_train_metric(
         self,
@@ -342,8 +342,7 @@ def update_train_metric(
         """
         if "x_hat" in predictions:
             embedding = self._extract_feature(batch)
-            mse = F.mse_loss(predictions["x_hat"], embedding, reduction="mean")
-            self._train_metric_modules["mse"].update(mse)
+            self._train_metric_modules["mse"].update(predictions["x_hat"], embedding)
 
     def update_metric(
         self,
@@ -360,7 +359,6 @@ def update_metric(
         """
         if "x_hat" in predictions:
             embedding = self._extract_feature(batch)
-            mse = F.mse_loss(predictions["x_hat"], embedding, reduction="mean")
-            self._metric_modules["mse"].update(mse)
+            self._metric_modules["mse"].update(predictions["x_hat"], embedding)
 
         self._update_unique_sid_ratio(predictions["codes"])

tzrec/modules/sid_generation/__init__.py

@@ -8,35 +8,3 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
-from tzrec.modules.sid_generation.kmeans import (
-    KMeansLayer,
-)
-from tzrec.modules.sid_generation.residual_kmeans_quantizer import (
-    ResidualKMeansQuantizer,
-)
-from tzrec.modules.sid_generation.residual_quantizer import (
-    ResidualQuantizer,
-)
-from tzrec.modules.sid_generation.residual_vector_quantizer import (
-    ResidualVectorQuantizer,
-)
-from tzrec.modules.sid_generation.types import (
-    QuantizeForwardMode,
-    QuantizeOutput,
-    ResidualQuantizerOutput,
-)
-from tzrec.modules.sid_generation.vector_quantize import (
-    VectorQuantize,
-)
-
-__all__ = [
-    "QuantizeForwardMode",
-    "QuantizeOutput",
-    "ResidualQuantizerOutput",
-    "VectorQuantize",
-    "ResidualQuantizer",
-    "ResidualVectorQuantizer",
-    "KMeansLayer",
-    "ResidualKMeansQuantizer",
-]