[feat] SID: add offline codebook collision-prevention tool

tzrec/models/sid_model.py

@@ -109,7 +109,7 @@ def _sid_predictions(
         """Build SID prediction tensors from a quantizer output."""
         predictions = {"codes": quant.cluster_ids}
         if quant.candidate_codes is not None and quant.candidate_scores is not None:
-            predictions["candidate_codes"] = quant.candidate_codes
+            predictions["candidate_codes"] = quant.candidate_codes.flatten(1)
             predictions["candidate_scores"] = quant.candidate_scores
         return predictions
 

tzrec/models/sid_rqkmeans_test.py

@@ -288,9 +288,9 @@ def test_inference_candidate_output_opt_in(self) -> None:
             set(preds.keys()), {"codes", "candidate_codes", "candidate_scores"}
         )
         self.assertEqual(preds["codes"].shape, (B, 2))
-        self.assertEqual(preds["candidate_codes"].shape, (B, 3, 2))
+        self.assertEqual(preds["candidate_codes"].shape, (B, 6))  # (B, topk*n_layers)
         self.assertEqual(preds["candidate_scores"].shape, (B, 3))
-        torch.testing.assert_close(preds["candidate_codes"][:, 0, :], preds["codes"])
+        torch.testing.assert_close(preds["candidate_codes"][:, :2], preds["codes"])
 
     def test_candidate_output_topk_exceeds_codebook_raises(self) -> None:
         """Reject topk above the last-layer codebook size at construction."""

tzrec/models/sid_rqvae_test.py

@@ -247,10 +247,12 @@ def test_rqvae_export_trace_keeps_candidate_output(self) -> None:
             {"codes", "candidate_codes", "candidate_scores"},
         )
         self.assertEqual(set(traced_predictions.keys()), set(predictions.keys()))
-        self.assertEqual(traced_predictions["candidate_codes"].shape, (B, 3, 2))
+        self.assertEqual(
+            traced_predictions["candidate_codes"].shape, (B, 6)
+        )  # (B, topk*n_layers)
         self.assertEqual(traced_predictions["candidate_scores"].shape, (B, 3))
         torch.testing.assert_close(
-            traced_predictions["candidate_codes"][:, 0, :],
+            traced_predictions["candidate_codes"][:, :2],
             traced_predictions["codes"],
         )
 
@@ -268,7 +270,7 @@ def test_rqvae_candidate_scores_sorted(self) -> None:
         self.assertEqual(
             set(preds.keys()), {"codes", "candidate_codes", "candidate_scores"}
         )
-        self.assertEqual(preds["candidate_codes"].shape, (B, 3, 2))
+        self.assertEqual(preds["candidate_codes"].shape, (B, 6))  # (B, topk*n_layers)
         self.assertEqual(preds["candidate_scores"].shape, (B, 3))
         scores = preds["candidate_scores"]
         # Slot 0 is the best-scored (nearest) neighbor: minimum score, sorted.