[feat] Add Qwen generative recommendation model

requirements/runtime.txt

@@ -22,3 +22,4 @@ tensorboard
 torch==2.12.1
 torchmetrics==1.0.3
 torchrec==1.7.0
+transformers @ https://tzrec.oss-accelerate.aliyuncs.com/third_party/transformers/transformers-4.51.2-py3-none-any.whl

tzrec/features/sid_feature.py

@@ -0,0 +1,174 @@
+# 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.
+
+from typing import Any, Dict, List, Optional, Tuple
+
+import numpy as np
+import pyarrow as pa
+
+from tzrec.datasets.utils import ParsedData
+from tzrec.features.feature import BaseFeature
+from tzrec.protos.feature_pb2 import FeatureConfig
+
+
+class SidFeature(BaseFeature):
+    """Semantic-ID sequence feature.
+
+    A flat stream of 0-based per-level SID codes -- whole items in level order --
+    plus the prompt text wrapping them. Under fg it is a passthrough.
+
+    Args:
+        feature_config (FeatureConfig): a instance of feature config.
+    """
+
+    def __init__(
+        self,
+        feature_config: FeatureConfig,
+        **kwargs: Any,
+    ) -> None:
+        # BaseFeature.__del__ dereferences _fg_op, so seed it before any raise.
+        self._fg_op = None
+        super().__init__(feature_config, **kwargs)
+        self._codebook = self._read_codebook()
+        # fg truncates by VALUE count and keeps the head, so a cap that is not a
+        # whole number of items would hand the model partial items. The model's
+        # own max_sequence_length is item-aligned and keeps the recent tail.
+        if self.config.HasField("sequence_length"):
+            if self.config.sequence_length % len(self._codebook):
+                raise ValueError(
+                    f"{self.__class__.__name__}[{self.config.feature_name}]: "
+                    f"sequence_length ({self.config.sequence_length}) must be a "
+                    f"multiple of the {len(self._codebook)}-level codebook, or "
+                    f"fg would cut an item in half. Prefer "
+                    f"model_config.common.max_sequence_length, which is "
+                    f"item-aligned and keeps the most RECENT items."
+                )
+        self._level_sizes = np.asarray(self._codebook)
+        self._level_offsets = np.cumsum(self._level_sizes) - self._level_sizes
+
+    def _read_codebook(self) -> List[int]:
+        """Validate the declared codebook once and normalize it to a list."""
+        codebook = [int(c) for c in self.config.codebook]
+        if not codebook:
+            raise ValueError(
+                f"{self.__class__.__name__}[{self.config.feature_name}]: codebook "
+                f"is required; give one vocabulary size per SID level."
+            )
+        if any(c <= 0 for c in codebook):
+            raise ValueError(
+                f"{self.__class__.__name__}[{self.config.feature_name}]: every "
+                f"codebook size must be positive, got {codebook}."
+            )
+        return codebook
+
+    @property
+    def value_dim(self) -> int:
+        """Fg value dimension of the feature."""
+        return self.config.value_dim
+
+    @property
+    def output_dim(self) -> int:
+        """Output dimension: SID codes pass through to the LM's own table."""
+        return self.value_dim
+
+    @property
+    def num_embeddings(self) -> int:
+        """Get embedding row count."""
+        raise RuntimeError(
+            f"{self.__class__.__name__}[{self.config.feature_name}] has no "
+            f"embedding table; SID codes index the LM vocabulary."
+        )
+
+    @property
+    def prefix_text(self) -> str:
+        """Text emitted immediately before this feature's SID tokens."""
+        return self.config.prefix_text
+
+    @property
+    def suffix_text(self) -> str:
+        """Text emitted immediately after this feature's SID tokens."""
+        return self.config.suffix_text
+
+    @property
+    def codebook(self) -> List[int]:
+        """Per-level SID vocabulary sizes; validated once at construction."""
+        return self._codebook
+
+    @property
+    def num_levels(self) -> int:
+        """Codes per item -- also the answer width."""
+        return len(self._codebook)
+
+    @property
+    def sid_vocab_size(self) -> int:
+        """Atoms the model must append to the backbone vocabulary."""
+        return sum(self._codebook)
+
+    @property
+    def level_offsets(self) -> List[int]:
+        """Flat offset of each level, i.e. ``cumsum(sizes) - sizes``."""
+        return self._level_offsets.tolist()
+
+    def _build_side_inputs(self) -> Optional[List[Tuple[str, str]]]:
+        """Input field names with side."""
+        if self.config.HasField("expression"):
+            return [tuple(self.config.expression.split(":"))]
+        else:
+            return None
+
+    def _parse(self, input_data: Dict[str, pa.Array]) -> ParsedData:
+        """Parse the SID stream into flat indices in the shared space.
+
+        Offsets are folded in here, in the dataloader workers: validating on the
+        forward path would let one rank raise and hang its peers on the
+        collective.
+        """
+        parsed = super()._parse(input_data)
+        num_levels = len(self._codebook)
+        bad = np.nonzero(parsed.seq_lengths % num_levels)[0]
+        if bad.size:
+            raise ValueError(
+                f"{self.__class__.__name__}[{self.config.feature_name}]: every "
+                f"row must hold whole {num_levels}-level items; rows "
+                f"{bad.tolist()[:10]} have lengths "
+                f"{parsed.seq_lengths[bad].tolist()[:10]}."
+            )
+        # rows are whole items, so each column is one level.
+        codes = parsed.values.reshape(-1, num_levels)
+        if ((codes < 0) | (codes >= self._level_sizes)).any():
+            raise ValueError(
+                f"{self.__class__.__name__}[{self.config.feature_name}]: SID "
+                f"codes must be local 0-based values in [0, codebook[level])."
+            )
+        # keep the dtype: int64 offsets would promote float32 to float64.
+        offsets = self._level_offsets.astype(codes.dtype, copy=False)
+        parsed.values = (codes + offsets).reshape(parsed.values.shape)
+        return parsed
+
+    def _fg_json(self) -> List[Dict[str, Any]]:
+        """Get fg json config impl.
+
+        A PASSTHROUGH: no fg feature_type can add ``level_offsets[i % levels]``,
+        so fg only reaches the codes and ``_parse`` does the arithmetic. It
+        exists because ``fg_mode`` is a data_config-level switch -- refusing it
+        here would block every other feature in the config.
+        """
+        # SCALAR form: fg_json prepends "sequence_" and injects the seq keys.
+        fg_cfg: Dict[str, Any] = {
+            "feature_type": "raw_feature",
+            "feature_name": self.config.feature_name,
+            "expression": self.config.expression,
+            "default_value": self.config.default_value,
+            "value_type": "float",
+        }
+        if self.config.HasField("stub_type"):
+            fg_cfg["stub_type"] = self.config.stub_type
+        return [fg_cfg]

tzrec/features/sid_feature_test.py

@@ -0,0 +1,136 @@
+# 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 pyarrow as pa
+from google.protobuf import text_format
+from parameterized import parameterized
+
+from tzrec.features.feature import FgMode, create_features
+from tzrec.protos import feature_pb2
+from tzrec.utils.test_util import parameterized_name_func
+
+
+def _feature(text, fg_mode=FgMode.FG_NONE):
+    fc = feature_pb2.FeatureConfig()
+    text_format.Merge(f"sequence_sid_feature {{ {text} }}", fc)
+    return create_features([fc], fg_mode=fg_mode)[0]
+
+
+_BASE = (
+    'feature_name: "user_sequence" expression: "user:user_sequence" '
+    "codebook: 4 codebook: 4 codebook: 4"
+)
+
+
+class SidFeatureTest(unittest.TestCase):
+    def test_dispatch_and_defaults(self) -> None:
+        f = _feature(_BASE)
+        self.assertEqual(type(f).__name__, "SidFeature")
+        # the oneof FIELD name is what makes it a sequence, not the message
+        self.assertTrue(f.is_sequence)
+        self.assertFalse(f.is_sparse)
+        self.assertEqual(f.name, "user_sequence")
+        self.assertEqual(f.value_dim, 1)
+        self.assertEqual(f.output_dim, 1)
+        self.assertEqual(f.side_inputs, [("user", "user_sequence")])
+        self.assertEqual(f.prefix_text, "")
+        self.assertEqual(f.suffix_text, "")
+        self.assertEqual(f.codebook, [4, 4, 4])
+        self.assertEqual(f.num_levels, 3)
+        self.assertEqual(f.sid_vocab_size, 12)
+        self.assertEqual(f.level_offsets, [0, 4, 8])
+
+    def test_prompt_text_round_trips(self) -> None:
+        f = _feature(f'{_BASE} prefix_text: "History: " suffix_text: "."')
+        self.assertEqual(f.prefix_text, "History: ")
+        self.assertEqual(f.suffix_text, ".")
+
+    @parameterized.expand(
+        [[FgMode.FG_NORMAL], [FgMode.FG_DAG], [FgMode.FG_BUCKETIZE]],
+        name_func=parameterized_name_func,
+    )
+    def test_builds_under_every_fg_mode(self, fg_mode) -> None:
+        # fg_mode is a data_config-level switch, so refusing it here would block
+        # every OTHER feature in the config from using fg.
+        self.assertEqual(type(_feature(_BASE, fg_mode=fg_mode)).__name__, "SidFeature")
+
+    def test_fg_passthrough_matches_the_fg_none_parse(self) -> None:
+        """Fg only reaches the codes; _parse folds the offsets either way."""
+        rows = [[1, 2, 3, 0, 1, 2], [3, 0, 1]]
+        want = [1, 6, 11, 0, 5, 10, 3, 4, 9]  # code + offsets [0,4,8]
+        none = _feature(_BASE).parse({"user_sequence": pa.array(rows)})
+        # under fg the same sequence arrives delimited, as ODPS/CSV deliver it
+        fg = _feature(_BASE, fg_mode=FgMode.FG_NORMAL).parse(
+            {"user_sequence": pa.array([";".join(map(str, r)) for r in rows])}
+        )
+        self.assertEqual(none.values.flatten().astype(int).tolist(), want)
+        self.assertEqual(fg.values.flatten().astype(int).tolist(), want)
+        self.assertEqual(none.seq_lengths.tolist(), fg.seq_lengths.tolist())
+
+    def test_fg_json_is_a_passthrough_raw_feature(self) -> None:
+        cfg = _feature(_BASE).fg_json()
+        self.assertEqual(len(cfg), 1)
+        # the base wrapper prepends "sequence_"; no bucketizer, no normalizer
+        self.assertEqual(cfg[0]["feature_type"], "sequence_raw_feature")
+        self.assertEqual(cfg[0]["expression"], "user:user_sequence")
+        for k in ("boundaries", "normalizer", "vocab_file", "hash_bucket_size"):
+            self.assertNotIn(k, cfg[0])
+
+    def test_rejects_a_sequence_length_that_splits_an_item(self) -> None:
+        # fg truncates by VALUE count, so a non-multiple would hand the model
+        # a partial item; _parse would then reject the whole batch.
+        with self.assertRaisesRegex(ValueError, "multiple of the 3-level"):
+            _feature(f"{_BASE} sequence_length: 10")
+        self.assertEqual(
+            _feature(f"{_BASE} sequence_length: 9").config.sequence_length, 9
+        )
+
+    def test_parse_folds_in_the_level_offsets(self) -> None:
+        # offsets [0, 4, 8]: level j's 0-based code k becomes flat index k + off[j],
+        # which is also the atom index -- no bridging shift anywhere.
+        f = _feature(_BASE)
+        parsed = f.parse({"user_sequence": pa.array([[0, 1, 2, 1, 2, 3], [0, 0, 0]])})
+        self.assertEqual(
+            parsed.values.flatten().tolist(), [0, 5, 10, 1, 6, 11, 0, 4, 8]
+        )
+        self.assertEqual(parsed.seq_lengths.tolist(), [6, 3])
+
+    def test_parse_rejects_out_of_range_and_partial_items(self) -> None:
+        f = _feature(_BASE)
+        with self.assertRaisesRegex(ValueError, "local 0-based"):
+            f.parse({"user_sequence": pa.array([[0, 1, 4]])})  # 4 == codebook[2]
+        with self.assertRaisesRegex(ValueError, "local 0-based"):
+            f.parse({"user_sequence": pa.array([[-1, 1, 2]])})
+        with self.assertRaisesRegex(ValueError, "whole 3-level items"):
+            f.parse({"user_sequence": pa.array([[0, 1]])})
+
+    def test_rejects_a_bad_codebook(self) -> None:
+        for bad, msg in (
+            ("", "codebook is required"),
+            ("codebook: 4 codebook: 0", "positive"),
+        ):
+            with self.subTest(bad=bad):
+                base = 'feature_name: "s" expression: "user:s" ' + bad
+                with self.assertRaisesRegex(ValueError, msg):
+                    _feature(base)
+
+    def test_no_embedding_table(self) -> None:
+        f = _feature(_BASE)
+        self.assertFalse(f.has_embedding)
+        self.assertIsNone(f.emb_config)
+        with self.assertRaisesRegex(RuntimeError, "no .*embedding table"):
+            _ = f.num_embeddings
+
+
+if __name__ == "__main__":
+    unittest.main()

tzrec/main.py

@@ -73,6 +73,7 @@
 from tzrec.optim.ema import DenseEMA, EMAOptimizer
 from tzrec.optim.lr_scheduler import BaseLR
 from tzrec.optim.optimizer import TZRecOptimizer
+from tzrec.protos import export_pb2
 from tzrec.protos.data_pb2 import DataConfig, DatasetType
 from tzrec.protos.eval_pb2 import EvalConfig
 from tzrec.protos.export_pb2 import ExportConfig
@@ -770,6 +771,9 @@ def train_and_evaluate(
         sample_weights=list(data_config.sample_weight_fields),
         sampler_type=sampler_type,
     )
+    # Cold start only; a resumed or fine-tuned run gets its weights from DCP.
+    if ckpt_path is None:
+        model.init_from_pretrained()
     model = TrainWrapper(
         model, device=device, mixed_precision=train_config.mixed_precision
     )
@@ -1073,6 +1077,35 @@ def export(
     if asset_files:
         assets = asset_files.split(",")
 
+    ckpt_manager = checkpoint_util.CheckpointManager(
+        pipeline_config.model_dir, export_config=pipeline_config.export_config
+    )
+    if not checkpoint_path:
+        if (
+            pipeline_config.HasField("export_config")
+            and pipeline_config.export_config.exporter_type == "best"
+        ):
+            checkpoint_path, _ = ckpt_manager.best_checkpoint()
+        else:
+            checkpoint_path, _ = ckpt_manager.latest_checkpoint()
+
+    # HF export converts the checkpoint dir directly -- no model build, no DCP restore.
+    if pipeline_config.export_config.export_format == export_pb2.ExportFormat.HF:
+        if not checkpoint_path:
+            raise ValueError("HF export: no checkpoint found to convert.")
+        if not os.path.exists(os.path.join(checkpoint_path, "config.json")):
+            raise ValueError(
+                f"HF export: {checkpoint_path} has no co-located HF assets; it "
+                f"was not written by an HF-backed model."
+            )
+        if assets:
+            logger.warning(f"HF export ignores asset_files: {assets}.")
+        if is_rank_zero:
+            from tzrec.utils.hf_export_util import dcp_to_hf
+
+            dcp_to_hf(checkpoint_path, export_dir)
+        return
+
     data_config = pipeline_config.data_config
 
     # Build feature
@@ -1092,18 +1125,6 @@ def export(
     model.set_is_inference(True)
     model = InferWrapper(model)
 
-    if not checkpoint_path:
-        ckpt_manager = checkpoint_util.CheckpointManager(
-            pipeline_config.model_dir, export_config=pipeline_config.export_config
-        )
-        if (
-            pipeline_config.HasField("export_config")
-            and pipeline_config.export_config.exporter_type == "best"
-        ):
-            checkpoint_path, _ = ckpt_manager.best_checkpoint()
-        else:
-            checkpoint_path, _ = ckpt_manager.latest_checkpoint()
-
     if isinstance(model.model, MatchModel):
         for name, module in model.model.named_children():
             if isinstance(module, MatchTower) or isinstance(module, MatchTowerWoEG):