[feat] prompt-native generative recommendation

requirements/runtime.txt

@@ -23,3 +23,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/datasets/dataset.py

@@ -18,6 +18,7 @@
 import numpy as np
 import pyarrow as pa
 import pyarrow.compute as pc
+import torch
 from torch import distributed as dist
 from torch.utils.data import DataLoader, IterableDataset, get_worker_info
 
@@ -40,6 +41,8 @@
     remove_nullable,
 )
 from tzrec.features.feature import BaseFeature
+from tzrec.prompt.assembler import PromptAssembler
+from tzrec.prompt.plan import CompiledPrompt, SlotSeg
 from tzrec.protos import data_pb2
 from tzrec.utils import config_util
 from tzrec.utils.load_class import get_register_class_meta
@@ -97,6 +100,7 @@ class BaseDataset(IterableDataset, metaclass=_dataset_meta_cls):
         mode (Mode): train or eval or predict.
         debug_level (int): dataset debug level, when mode=predict and
             debug_level > 0, will dump fg encoded data to debug_str
+        compiled_prompt (CompiledPrompt, optional): compiled prompt assembly contract.
     """
 
     def __init__(
@@ -107,8 +111,18 @@ def __init__(
         reserved_columns: Optional[List[str]] = None,
         mode: Mode = Mode.EVAL,
         debug_level: int = 0,
+        compiled_prompt: Optional[CompiledPrompt] = None,
     ) -> None:
         super(BaseDataset, self).__init__()
+        self._prompt_assembler = (
+            PromptAssembler(
+                compiled_prompt.prompt_plan,
+                compiled_prompt.sid_space,
+                include_response=mode != Mode.PREDICT,
+            )
+            if compiled_prompt is not None
+            else None
+        )
         self._data_config = data_config
         self._features = features
         self._input_path = input_path
@@ -121,8 +135,25 @@ def __init__(
             else None
         )
 
+        parser_features = features
+        if compiled_prompt is not None and mode == Mode.PREDICT:
+            prompt_feature_names = {
+                feature_name
+                for segment in compiled_prompt.prompt_plan.segments
+                if isinstance(segment, SlotSeg)
+                for feature_name in segment.feature_names
+            }
+            response_feature_names = {
+                feature_name
+                for segment in compiled_prompt.prompt_plan.response_segments
+                if isinstance(segment, SlotSeg)
+                for feature_name in segment.feature_names
+            }
+            response_only = response_feature_names - prompt_feature_names
+            parser_features = [f for f in features if f.name not in response_only]
+
         self._data_parser = DataParser(
-            features=features,
+            features=parser_features,
             labels=list(data_config.label_fields)
             if self._mode != Mode.PREDICT
             else None,
@@ -385,6 +416,16 @@ def _build_batch(self, input_data: Dict[str, pa.Array]) -> Batch:
         else:
             batch = self._data_parser.to_batch(output_data)
 
+        if self._prompt_assembler is not None:
+            batch.additional_infos.update(
+                {
+                    k: torch.from_numpy(np.asarray(v))
+                    for k, v in self._prompt_assembler.assemble_batch(
+                        output_data
+                    ).items()
+                }
+            )
+
         # Set checkpoint info on batch
         batch.checkpoint_info = checkpoint_info
         batch.data_timestamp = data_timestamp
@@ -762,6 +803,7 @@ def create_dataloader(
     gl_cluster: Optional[Dict[str, Union[int, str]]] = None,
     debug_level: int = 0,
     checkpoint_state: Optional[Dict[str, Any]] = None,
+    compiled_prompt: Optional[CompiledPrompt] = None,
 ) -> DataLoader:
     """Build dataloader.
 
@@ -776,6 +818,8 @@ def create_dataloader(
             debug_level > 0, will dump fg encoded data to debug_str
         checkpoint_state (dict, optional): resume state, applied before the
             eager ``iter()`` forks workers so it reaches them.
+        compiled_prompt (CompiledPrompt, optional): when set, each batch carries
+            the assembled prompt streams in ``additional_infos``.
 
     Return:
         dataloader (dataloader): a DataLoader.
@@ -790,6 +834,7 @@ def create_dataloader(
         reserved_columns=reserved_columns,
         mode=mode,
         debug_level=debug_level,
+        compiled_prompt=compiled_prompt,
     )
     if checkpoint_state:
         dataset.load_state_dict(dict(checkpoint_state))

tzrec/main.py

@@ -72,12 +72,17 @@
 from tzrec.optim.ema import DenseEMA, EMAOptimizer
 from tzrec.optim.lr_scheduler import BaseLR
 from tzrec.optim.optimizer import TZRecOptimizer
+from tzrec.prompt.compile import compile_prompt
+from tzrec.prompt.persist import check_prompt_assets, copy_prompt_assets
+from tzrec.prompt.plan import CompiledPrompt
+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
 from tzrec.protos.feature_pb2 import FeatureConfig
 from tzrec.protos.model_pb2 import Kernel as KernelProto
 from tzrec.protos.model_pb2 import ModelConfig
+from tzrec.protos.pipeline_pb2 import EasyRecConfig
 from tzrec.protos.train_pb2 import TrainConfig
 from tzrec.utils import checkpoint_util, config_util, predict_util
 from tzrec.utils.delta_embedding_dump import DeltaEmbeddingDumper
@@ -119,6 +124,17 @@ def _create_features(
     return features
 
 
+def _compile_prompt(
+    pipeline_config: EasyRecConfig, features: List[BaseFeature]
+) -> Optional[CompiledPrompt]:
+    """Compile prompt_config for entry points that build prompt-aware objects."""
+    if not pipeline_config.HasField("prompt_config"):
+        return None
+    return compile_prompt(
+        pipeline_config.prompt_config, features, model_dir=pipeline_config.model_dir
+    )
+
+
 def _get_sampler_type(data_config: DataConfig) -> Optional[str]:
     try:
         sampler_type = (
@@ -137,6 +153,7 @@ def _create_model(
     labels: List[str],
     sample_weights: Optional[List[str]] = None,
     sampler_type: Optional[str] = None,
+    compiled_prompt: Optional[CompiledPrompt] = None,
 ) -> BaseModel:
     """Build model.
 
@@ -146,6 +163,8 @@ def _create_model(
         labels (list): list of label names.
         sample_weights (list): list of sample weight names.
         sampler_type (str): negative sampler type
+        compiled_prompt (CompiledPrompt, optional): forwarded to prompt-native models.
+
     Return:
         model: a EasyRec Model.
     """
@@ -159,6 +178,7 @@ def _create_model(
         labels,
         sample_weights=sample_weights,
         sampler_type=sampler_type,
+        compiled_prompt=compiled_prompt,
     )
 
     kernel = Kernel[KernelProto.Name(model_config.kernel)]
@@ -693,6 +713,7 @@ def train_and_evaluate(
     data_config = pipeline_config.data_config
     # Build feature
     features = _create_features(list(pipeline_config.feature_configs), data_config)
+    compiled_prompt = _compile_prompt(pipeline_config, features)
 
     ckpt_manager = checkpoint_util.CheckpointManager(
         pipeline_config.model_dir,
@@ -731,6 +752,8 @@ def train_and_evaluate(
 
     # Restore dataloader state before create_dataloader starts its workers
     dataloader_state: Optional[Dict[str, Any]] = None
+    if ckpt_path:
+        check_prompt_assets(compiled_prompt, ckpt_path)
     if ckpt_path and continue_train:
         dataloader_state = ckpt_manager.restore_dataloader_state(ckpt_path)
         if dataloader_state and not restore_from_model_dir:
@@ -743,6 +766,7 @@ def train_and_evaluate(
         features,
         pipeline_config.train_input_path,
         mode=Mode.TRAIN,
+        compiled_prompt=compiled_prompt,
         checkpoint_state=dataloader_state,
     )
     eval_dataloader = None
@@ -754,6 +778,7 @@ def train_and_evaluate(
             features,
             pipeline_config.eval_input_path,
             mode=Mode.EVAL,
+            compiled_prompt=compiled_prompt,
             gl_cluster=gl_cluster,
         )
 
@@ -766,7 +791,11 @@ def train_and_evaluate(
         list(data_config.label_fields),
         sample_weights=list(data_config.sample_weight_fields),
         sampler_type=sampler_type,
+        compiled_prompt=compiled_prompt,
     )
+    # 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
     )
@@ -957,12 +986,14 @@ def evaluate(
     data_config = pipeline_config.data_config
     # Build feature
     features = _create_features(list(pipeline_config.feature_configs), data_config)
+    compiled_prompt = _compile_prompt(pipeline_config, features)
 
     eval_dataloader = create_dataloader(
         data_config,
         features,
         eval_input_path or pipeline_config.eval_input_path,
         mode=Mode.EVAL,
+        compiled_prompt=compiled_prompt,
     )
 
     sampler_type = _get_sampler_type(data_config)
@@ -974,6 +1005,7 @@ def evaluate(
         list(data_config.label_fields),
         sample_weights=list(data_config.sample_weight_fields),
         sampler_type=sampler_type,
+        compiled_prompt=compiled_prompt,
     )
     model = TrainWrapper(
         model, device=device, mixed_precision=train_config.mixed_precision
@@ -1005,6 +1037,7 @@ def evaluate(
     )
 
     if checkpoint_path:
+        check_prompt_assets(compiled_prompt, checkpoint_path)
         ckpt_manager.restore(
             checkpoint_path,
             model,
@@ -1077,6 +1110,53 @@ 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 config_util.use_dense_ema(
+            pipeline_config.export_config, pipeline_config.train_config
+        ):
+            raise ValueError(
+                "HF export: dcp_to_hf reads <checkpoint>/model, so it cannot "
+                "serve Dense EMA parameters. Set export_config.use_dense_ema to "
+                "false to export the raw weights."
+            )
+        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)
+            # Carry the prompt contract saved alongside the checkpoint.
+            copy_prompt_assets(checkpoint_path, export_dir)
+        return
+
+    if pipeline_config.HasField("prompt_config"):
+        raise ValueError(
+            "a prompt-native model exports to a HuggingFace directory, not "
+            "TorchScript: its input is an assembled token stream the dataloader "
+            "builds, which an export-time dummy batch cannot supply. Set "
+            "export_config.export_format to HF."
+        )
+
     data_config = pipeline_config.data_config
 
     # Build feature
@@ -1096,18 +1176,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):
@@ -1284,13 +1352,15 @@ def predict(
     data_config.drop_remainder = False
     # Build feature
     features = _create_features(list(pipeline_config.feature_configs), data_config)
+    compiled_prompt = _compile_prompt(pipeline_config, features)
 
     infer_dataloader = create_dataloader(
         data_config,
         features,
         predict_input_path,
         reserved_columns=reserved_cols,
         mode=Mode.PREDICT,
+        compiled_prompt=compiled_prompt,
         debug_level=debug_level,
     )
     infer_iterator = infer_dataloader.get_iterator()  # pyre-ignore[16]
@@ -1557,6 +1627,7 @@ def predict_checkpoint(
     data_config = pipeline_config.data_config
     # Build feature
     features = _create_features(list(pipeline_config.feature_configs), data_config)
+    compiled_prompt = _compile_prompt(pipeline_config, features)
 
     # Build dataloader
     predict_dataloader = create_dataloader(
@@ -1565,6 +1636,7 @@ def predict_checkpoint(
         predict_input_path,
         reserved_columns=reserved_cols,
         mode=Mode.PREDICT,
+        compiled_prompt=compiled_prompt,
         debug_level=debug_level,
     )
 
@@ -1584,6 +1656,7 @@ def predict_checkpoint(
         pipeline_config.model_config,
         features,
         [],
+        compiled_prompt=compiled_prompt,
     )
     model.set_is_inference(True)
     model = PredictWrapper(
@@ -1619,6 +1692,7 @@ def predict_checkpoint(
     model.eval()
 
     if checkpoint_path:
+        check_prompt_assets(compiled_prompt, checkpoint_path)
         ckpt_manager.restore(
             checkpoint_path,
             model,