[review] SID base: second round of github-actions feedback on PR alibaba#538

- ResidualQuantizer.__init__: assert n_layers >= 1, making the boundary
  self-guarding (matches the "n_layers >= 1 is guaranteed" comment in
  decode_codes; a future n_layers=0 subclass would otherwise index OOB).
- BaseSidModel.update_train_metric: type predictions as Dict[str, Tensor]
  instead of bare dict, matching BaseModel / sibling models.
- sid_model_test.py (new): unit-test _update_unique_sid_ratio via __new__ —
  empty-batch no-op (guard) and the exact 0.75 ratio on a known-duplicate
  batch.
- residual_quantizer_test: make the normalize_residuals test behavioral —
  toggle the flag on the same input/codebook and assert the assignments
  differ, so the normalization branch is actually validated (the old shape
  assertions passed even with the branch deleted).

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

Author: WhiteSwan1

tzrec/models/sid_model.py

@@ -11,7 +11,7 @@
 
 """BaseSidModel: shared base for semantic-ID generation models."""
 
-from typing import Any, List, Optional
+from typing import Any, Dict, List, Optional
 
 import torch
 import torchmetrics
@@ -105,7 +105,7 @@ def init_metric(self) -> None:
 
     def update_train_metric(
         self,
-        predictions: dict,
+        predictions: Dict[str, torch.Tensor],
         batch: Batch,
     ) -> None:
         """Update train-path metric state.

tzrec/models/sid_model_test.py

@@ -0,0 +1,55 @@
+# 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
+import torchmetrics
+
+from tzrec.models.sid_model import BaseSidModel
+
+
+class UpdateUniqueSidRatioTest(unittest.TestCase):
+    """Unit-test the codebook-coverage metric helper.
+
+    ``_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.
+    """
+
+    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()}
+        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
+        )
+
+    def test_ratio_on_known_duplicates(self) -> None:
+        model = self._bare_model()
+        # 3 unique rows out of 4 -> ratio 0.75.
+        codes = torch.tensor([[1, 2], [1, 2], [3, 4], [5, 6]])
+        model._update_unique_sid_ratio(codes)
+        self.assertAlmostEqual(
+            model._metric_modules["unique_sid_ratio"].compute().item(), 0.75, places=6
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

tzrec/modules/sid_generation/residual_quantizer.py

@@ -78,6 +78,7 @@ def __init__(
         normalize_residuals: bool = False,
     ) -> None:
         super().__init__()
+        assert n_layers >= 1, f"n_layers must be >= 1, got {n_layers}"
         self.embed_dim = embed_dim
         self.n_layers = n_layers
         self.normalize_residuals = normalize_residuals

tzrec/modules/sid_generation/residual_quantizer_test.py

@@ -119,14 +119,21 @@ def test_detach_invariant(self) -> None:
         self.assertIsNotNone(fq.books[0].grad)
         self.assertIsNone(x.grad)
 
-    def test_normalize_residuals_branch(self) -> None:
-        torch.manual_seed(0)
-        fq = _FakeQuantizer(
-            embed_dim=4, n_layers=2, n_embed=5, normalize_residuals=True
+    def test_normalize_residuals_changes_assignment(self) -> None:
+        # Same input and same codebook (re-seeded before each build), so the
+        # only difference is the normalize_residuals branch — it must change
+        # the residual a later layer sees and hence the codes it assigns.
+        x = torch.randn(8, 4)
+        torch.manual_seed(1)
+        fq_off = _FakeQuantizer(embed_dim=4, n_layers=2, n_embed=6)
+        torch.manual_seed(1)
+        fq_on = _FakeQuantizer(
+            embed_dim=4, n_layers=2, n_embed=6, normalize_residuals=True
         )
-        ids, agg, _ = fq._residual_pass(torch.randn(5, 4))
-        self.assertEqual(ids.shape, (5, 2))
-        self.assertEqual(agg.shape, (5, 4))
+        ids_off, _, _ = fq_off._residual_pass(x)
+        ids_on, _, _ = fq_on._residual_pass(x)
+        self.assertEqual(ids_on.shape, (8, 2))
+        self.assertFalse(torch.equal(ids_off, ids_on))
 
     def test_decode_codes_sum_and_dtype(self) -> None:
         torch.manual_seed(0)