Refactor training transforms to reuse the inference transforms.

PiperOrigin-RevId: 869786054

Author: The kauldron Authors

kauldron/contrib/data/__init__.py

@@ -72,6 +72,7 @@
   from kauldron.contrib.data.preprocessing import TemporalRandomWalk
   from kauldron.contrib.data.preprocessing import TemporalRandomWindow
   from kauldron.contrib.data.preprocessing import TimeChunkedFlattenVideo
+  from kauldron.contrib.data.preprocessing import TreeUnflattenForKey
   from kauldron.contrib.data.preprocessing import VStack
   from kauldron.contrib.data.preprocessing import ValueRange
   from kauldron.contrib.data.preprocessing import VideoMAENormalization

kauldron/contrib/data/preprocessing.py

@@ -19,6 +19,7 @@
 import abc
 from collections.abc import Mapping
 import dataclasses
+import functools
 import typing
 from typing import Any, Callable, List, Literal, Optional, Sequence
 
@@ -1630,3 +1631,63 @@ def map(self, features: dict[str, tf.Tensor]) -> dict[str, tf.Tensor]:
 
     features[self.output_key] = final_video
     return features
+
+
+@dataclasses.dataclass(kw_only=True, frozen=True, eq=True)
+class TreeUnflattenForKey(kd.data.MapTransform):
+  """Unflattens a previously flattened dictionary within an element.
+
+  This transform is designed to reverse the effect of
+  kd.data.TreeFlattenWithPath.
+  """
+
+  key: str
+  separator: str = "_"
+
+  @functools.cached_property
+  def _prefix(self) -> str:
+    return self.key + self.separator
+
+  def map(self, element: dict[str, Any]) -> dict[str, Any]:
+    element = dict(element)  # Ensure the element is mutable
+    flat_subtree = {}
+    keys_to_remove = []
+
+    # Extract keys belonging to the flattened subtree
+    for k, v in element.items():
+      if k.startswith(self._prefix):
+        flat_subtree[k] = v
+        keys_to_remove.append(k)
+
+    if not flat_subtree:
+      # No keys to unflatten for the specified key
+      return element
+
+    # Reconstruct the nested dictionary
+    nested_subtree = {}
+    for long_key, value in flat_subtree.items():
+      # Remove the prefix
+      path_str = long_key[len(self._prefix) :]
+      parts = path_str.split(self.separator)
+
+      current_level = nested_subtree
+      for i, part in enumerate(parts):
+        if i == len(parts) - 1:
+          # Last part of the path, assign the value
+          current_level[part] = value
+        else:
+          # Navigate/create nested dictionaries
+          if part not in current_level:
+            current_level[part] = {}
+          elif not isinstance(current_level[part], dict):
+            # This case should ideally not happen if flattening was consistent
+            # Handle potential conflicts if a path is both a leaf and a branch
+            raise ValueError(f"Conflict at path {'.'.join(parts[:i+1])}")
+          current_level = current_level[part]
+
+    # Update the element: add the nested structure, remove old flat keys
+    element[self.key] = nested_subtree
+    for k in keys_to_remove:
+      del element[k]
+
+    return element