[transform] Add conversation trace hooks

This adds minimal postprocessing and row-ID hook seams for SFT conversation transforms so Hermes-style trace normalization can land without reshaping the existing adapter stack. It also preserves transformed dataset hash stability when the new hooks are unset and adds regression tests for hook ordering and signature behavior.

Author: taivu1998

experiments/posttrain/instruction_datasets.py

@@ -37,7 +37,6 @@
 22. open-thoughts/OpenThoughts3-1.2M  # Original OT3 dataset; smoltalk2 uses a slightly different version
 """
 
-import dataclasses
 import hashlib
 import json
 from collections.abc import Sequence
@@ -57,7 +56,13 @@
     ConversationToDolmaConfig,
     convert_conversation_to_dolma,
 )
-from marin.transform.conversation.adapters import InputDatasetFormat, TransformAdapter
+from marin.transform.conversation.adapters import (
+    InputDatasetFormat,
+    MessagePostprocessFn,
+    RowIdFn,
+    TransformAdapter,
+    transform_adapter_signature,
+)
 from marin.transform.conversation.transform_conversation import (
     TransformSFTDatasetConfig,
     transform_hf_dataset,
@@ -142,6 +147,8 @@ def multi_turn_adapter(
     metadata_remap: dict[str, str] | None = None,
     replacements: dict[str, str] | None = None,
     extra_metadata_fn=None,
+    message_postprocess_fn: MessagePostprocessFn | None = None,
+    row_id_fn: RowIdFn | None = None,
 ) -> TransformAdapter:
     return TransformAdapter(
         dataset_format=InputDatasetFormat.SINGLE_COLUMN_MULTI_TURN,
@@ -154,6 +161,8 @@ def multi_turn_adapter(
         metadata_remap=metadata_remap or {},
         replacements=replacements,
         extra_metadata_fn=extra_metadata_fn,
+        message_postprocess_fn=message_postprocess_fn,
+        row_id_fn=row_id_fn,
     )
 
 
@@ -166,6 +175,8 @@ def instruction_response_adapter(
     metadata_remap: dict[str, str] | None = None,
     replacements: dict[str, str] | None = None,
     extra_metadata_fn=None,
+    message_postprocess_fn: MessagePostprocessFn | None = None,
+    row_id_fn: RowIdFn | None = None,
 ) -> TransformAdapter:
     return TransformAdapter(
         dataset_format=InputDatasetFormat.INSTRUCTION_RESPONSE,
@@ -176,6 +187,8 @@ def instruction_response_adapter(
         metadata_remap=metadata_remap or {},
         replacements=replacements,
         extra_metadata_fn=extra_metadata_fn,
+        message_postprocess_fn=message_postprocess_fn,
+        row_id_fn=row_id_fn,
     )
 
 
@@ -186,6 +199,8 @@ def instruct_column_response_adapter(
     metadata_remap: dict[str, str] | None = None,
     replacements: dict[str, str] | None = None,
     extra_metadata_fn=None,
+    message_postprocess_fn: MessagePostprocessFn | None = None,
+    row_id_fn: RowIdFn | None = None,
 ) -> TransformAdapter:
     return TransformAdapter(
         dataset_format=InputDatasetFormat.INSTRUCT_COLUMN_RESPONSE,
@@ -195,6 +210,8 @@ def instruct_column_response_adapter(
         metadata_remap=metadata_remap or {},
         replacements=replacements,
         extra_metadata_fn=extra_metadata_fn,
+        message_postprocess_fn=message_postprocess_fn,
+        row_id_fn=row_id_fn,
     )
 
 
@@ -210,6 +227,8 @@ def instruct_msg_response_adapter(
     metadata_remap: dict[str, str] | None = None,
     replacements: dict[str, str] | None = None,
     extra_metadata_fn=None,
+    message_postprocess_fn: MessagePostprocessFn | None = None,
+    row_id_fn: RowIdFn | None = None,
 ) -> TransformAdapter:
     return TransformAdapter(
         dataset_format=InputDatasetFormat.INSTRUCT_MSG_RESPONSE,
@@ -223,6 +242,8 @@ def instruct_msg_response_adapter(
         metadata_remap=metadata_remap or {},
         replacements=replacements,
         extra_metadata_fn=extra_metadata_fn,
+        message_postprocess_fn=message_postprocess_fn,
+        row_id_fn=row_id_fn,
     )
 
 
@@ -561,26 +582,18 @@ def get_directory_friendly_dataset_name(hf_dataset_id: str) -> str:
     return dataset_name
 
 
+def get_adapter_signature_string(adapter: TransformAdapter) -> str:
+    """Return the stable JSON signature used to version transformed instruction datasets."""
+    return json.dumps(transform_adapter_signature(adapter), sort_keys=True)
+
+
 def transform_dataset_step(dataset_cfg: InstructionDatasetConfig) -> ExecutorStep:
     """ExecutorStep that preprocesses the input dataset into a canonicalized format for SFT training."""
     adapter = dataset_cfg.adapter
     output_name = dataset_cfg.name if dataset_cfg.name is not None else dataset_cfg.hf_dataset_id
     dataset_name = get_directory_friendly_dataset_name(output_name)
 
-    adapter_dict = dataclasses.asdict(adapter)
-    adapter_dict["dataset_format"] = adapter_dict["dataset_format"].value
-
-    def canonicalize(value):
-        if isinstance(value, dict):
-            return {k: canonicalize(v) for k, v in sorted(value.items())}
-        if isinstance(value, list):
-            return [canonicalize(x) for x in value]
-        if callable(value):
-            return f"{value.__module__}.{value.__qualname__}"
-        return value
-
-    adapter_signature = canonicalize(adapter_dict)
-    adapter_signature_str = json.dumps(adapter_signature, sort_keys=True)
+    adapter_signature_str = get_adapter_signature_string(adapter)
 
     config_str = f"{dataset_name}-\
         {dataset_cfg.revision}\

lib/marin/src/marin/transform/conversation/adapters.py

@@ -55,6 +55,36 @@ class InputDatasetFormat(str, Enum):
     INSTRUCT_MSG_RESPONSE: str = "instruct_msg_response"
 
 
+MessagePostprocessFn = Callable[[list[OpenAIChatMessage], dict[str, Any]], list[OpenAIChatMessage]]
+RowIdFn = Callable[[dict[str, Any], list[dict[str, Any]]], str]
+
+_OPTIONAL_SIGNATURE_FIELDS = frozenset({"message_postprocess_fn", "row_id_fn"})
+
+
+def _canonicalize_signature_value(value: Any) -> Any:
+    if isinstance(value, dict):
+        return {k: _canonicalize_signature_value(v) for k, v in sorted(value.items())}
+    if isinstance(value, list):
+        return [_canonicalize_signature_value(x) for x in value]
+    if callable(value):
+        return f"{value.__module__}.{value.__qualname__}"
+    return value
+
+
+def transform_adapter_signature(adapter: "TransformAdapter") -> dict[str, Any]:
+    """Return a stable, JSON-serializable signature for a transform adapter.
+
+    Newly added optional trace hooks are omitted when unset so existing dataset output hashes
+    remain stable.
+    """
+    adapter_dict = dataclasses.asdict(adapter)
+    adapter_dict["dataset_format"] = adapter_dict["dataset_format"].value
+    adapter_dict = {
+        key: value for key, value in adapter_dict.items() if not (key in _OPTIONAL_SIGNATURE_FIELDS and value is None)
+    }
+    return _canonicalize_signature_value(adapter_dict)
+
+
 @dataclass
 class TransformAdapter:
     dataset_format: InputDatasetFormat = InputDatasetFormat.INSTRUCTION_RESPONSE
@@ -87,6 +117,8 @@ class TransformAdapter:
     metadata_remap: dict[str, str] = field(default_factory=dict)
     replacements: dict[str, str] | None = None
     extra_metadata_fn: Callable[[dict[str, Any]], dict[str, Any]] | None = None
+    message_postprocess_fn: MessagePostprocessFn | None = None
+    row_id_fn: RowIdFn | None = None
 
     def transform_conversation_to_openai_format(
         self,

lib/marin/src/marin/transform/conversation/transform_conversation.py

@@ -126,10 +126,10 @@ def transform_row(row: dict, cfg: TransformSFTDatasetConfig, adapter: TransformA
         logger.warning(f"{source} returning no valid messages")
         return None
 
-    transformed_row_messages = [message.model_dump() for message in transformed_row_messages]
+    if adapter.message_postprocess_fn:
+        transformed_row_messages = adapter.message_postprocess_fn(transformed_row_messages, row)
 
-    # Create a unique ID for the row based on the text
-    row_idx = generate_hash_from_messages(transformed_row_messages)
+    transformed_row_messages = [message.model_dump() for message in transformed_row_messages]
     metadata_columns = unwrap_versioned_value(cfg.metadata_columns)
     metadata_remap = adapter.metadata_remap or {}
     replacements = adapter.replacements if adapter.replacements is not None else DEFAULT_TEXT_REPLACEMENTS
@@ -154,6 +154,13 @@ def transform_row(row: dict, cfg: TransformSFTDatasetConfig, adapter: TransformA
         transformed_row_messages = [_normalize_tool_structures(message) for message in transformed_row_messages]
     else:
         transformed_row_messages = [_normalize_tool_structures(message) for message in transformed_row_messages]
+
+    if adapter.row_id_fn:
+        row_idx = adapter.row_id_fn(row, transformed_row_messages)
+    else:
+        # Create a unique ID for the row based on the transformed text.
+        row_idx = generate_hash_from_messages(transformed_row_messages)
+
     if adapter.extra_metadata_fn:
         extra_from_fn = adapter.extra_metadata_fn(row)
         if extra_from_fn:

tests/transform/test_conversation.py

@@ -3,8 +3,12 @@
 
 """Tests for conversation data transformation scripts."""
 
+import dataclasses
+import json
 from pathlib import Path
 
+from experiments.posttrain.instruction_datasets import get_adapter_signature_string, instruction_response_adapter
+from marin.core.conversation import OpenAIChatMessage
 from marin.transform.conversation.adapters import InputDatasetFormat, TransformAdapter
 from marin.transform.conversation.conversation_to_dolma import transform_conversation_to_dolma
 from marin.transform.conversation.preference_data_adapters import PreferenceTransformAdapter
@@ -49,6 +53,35 @@
 }
 
 
+def _replace_assistant_message(messages: list[OpenAIChatMessage], row: dict[str, str]) -> list[OpenAIChatMessage]:
+    return [
+        messages[0],
+        OpenAIChatMessage(role="assistant", content=row["postprocessed_response"]),
+    ]
+
+
+def _row_id_from_source(row: dict[str, str], _messages: list[dict[str, object]]) -> str:
+    return row["custom_row_id"]
+
+
+def _legacy_adapter_signature_string(adapter: TransformAdapter) -> str:
+    adapter_dict = dataclasses.asdict(adapter)
+    adapter_dict["dataset_format"] = adapter_dict["dataset_format"].value
+    adapter_dict.pop("message_postprocess_fn", None)
+    adapter_dict.pop("row_id_fn", None)
+
+    def canonicalize(value):
+        if isinstance(value, dict):
+            return {key: canonicalize(inner_value) for key, inner_value in sorted(value.items())}
+        if isinstance(value, list):
+            return [canonicalize(item) for item in value]
+        if callable(value):
+            return f"{value.__module__}.{value.__qualname__}"
+        return value
+
+    return json.dumps(canonicalize(adapter_dict), sort_keys=True)
+
+
 class TestTransformAdapters:
     """Test the different adapter formats."""
 
@@ -127,6 +160,94 @@ def test_transform_with_replacements(self):
         assert "<|end_think|>" in response_message.content
         assert "<think>" not in response_message.content
 
+    def test_transform_applies_message_postprocess_before_replacements(self):
+        """Test message postprocessing runs before text replacements."""
+        adapter = TransformAdapter(
+            dataset_format=InputDatasetFormat.INSTRUCTION_RESPONSE,
+            instruction_column="instruction",
+            response_column="response",
+            replacements={"<think>": "<|start_think|>", "</think>": "<|end_think|>"},
+            message_postprocess_fn=_replace_assistant_message,
+        )
+
+        row = {
+            "instruction": "Solve this",
+            "response": "placeholder",
+            "postprocessed_response": "<think>Use the replacement path</think>",
+        }
+
+        cfg = TransformSFTDatasetConfig(
+            source="test/dataset",
+            revision="main",
+            output_path="/tmp/output",
+            metadata_columns=[],
+            adapter=adapter,
+        )
+
+        result = transform_row(row, cfg, adapter)
+
+        assert result is not None
+        response_message = result.messages[1]
+        assert response_message.content == "<|start_think|>Use the replacement path<|end_think|>"
+
+    def test_transform_uses_row_id_hook(self):
+        """Test row ids can come from a source-provided identifier."""
+        adapter = TransformAdapter(
+            dataset_format=InputDatasetFormat.INSTRUCTION_RESPONSE,
+            instruction_column="instruction",
+            response_column="response",
+            row_id_fn=_row_id_from_source,
+        )
+
+        row = {
+            "instruction": "Question",
+            "response": "Answer",
+            "custom_row_id": "trace-123",
+        }
+
+        cfg = TransformSFTDatasetConfig(
+            source="test/dataset",
+            revision="main",
+            output_path="/tmp/output",
+            metadata_columns=[],
+            adapter=adapter,
+        )
+
+        result = transform_row(row, cfg, adapter)
+
+        assert result is not None
+        assert result.id == "trace-123"
+
+
+class TestInstructionDatasetAdapterSignatures:
+    """Test instruction dataset adapter signature stability."""
+
+    def test_signature_omits_unset_trace_hooks(self):
+        adapter = instruction_response_adapter(
+            instruction_column="instruction",
+            response_column="response",
+        )
+
+        signature_string = get_adapter_signature_string(adapter)
+        signature = json.loads(signature_string)
+
+        assert "message_postprocess_fn" not in signature
+        assert "row_id_fn" not in signature
+        assert signature_string == _legacy_adapter_signature_string(adapter)
+
+    def test_signature_includes_set_trace_hooks(self):
+        adapter = instruction_response_adapter(
+            instruction_column="instruction",
+            response_column="response",
+            message_postprocess_fn=_replace_assistant_message,
+            row_id_fn=_row_id_from_source,
+        )
+
+        signature = json.loads(get_adapter_signature_string(adapter))
+
+        assert signature["message_postprocess_fn"].endswith("._replace_assistant_message")
+        assert signature["row_id_fn"].endswith("._row_id_from_source")
+
 
 class TestPreferenceDataTransform:
     """Test preference data (DPO) transformation."""