code-loader 1.0.191.dev1__py3-none-any.whl → 1.0.193__py3-none-any.whl

code_loader/contract/datasetclasses.py

@@ -1,6 +1,6 @@
 import warnings
 from dataclasses import dataclass, field
-from typing import Any, Callable, List, Optional, Dict, Union, Type, Literal, cast
+from typing import Any, Callable, List, Optional, Dict, Union, Type, Literal
 import re
 import numpy as np
 import numpy.typing as npt
@@ -15,8 +15,6 @@ custom_latent_space_attribute = "custom_latent_space"
 
 _simulation_context: Dict[str, bool] = {"active": False}
 
-SampleId = Union[int, str]
-
 
 @dataclass
 class PreprocessResponse:
@@ -41,44 +39,27 @@ class PreprocessResponse:
     """
     length: Optional[int] = None  # Deprecated. Please use sample_ids instead
     data: Any = None
-    sample_ids: Optional[Union[List[SampleId], List[List[SampleId]]]] = None
+    sample_ids: Optional[Union[List[str], List[int]]] = None
     state: Optional[DataStateType] = None
     sample_id_type: Optional[Union[Type[str], Type[int]]] = None
     sample_ids_to_instance_mappings: Optional[Dict[str, List[str]]] = None # in use only for element instance
     instance_to_sample_ids_mappings: Optional[Dict[str, str]] = None # in use only for element instance
     tl_generated: bool = False
-    _grouped: bool = field(default=False, init=False, repr=False, compare=False)
 
     def __post_init__(self) -> None:
         assert self.sample_ids_to_instance_mappings is None, f"Keep sample_ids_to_instance_mappings None when initializing PreprocessResponse"
         assert self.instance_to_sample_ids_mappings is None, f"Keep instance_to_sample_ids_mappings None when initializing PreprocessResponse"
 
         if self.length is not None and self.sample_ids is None:
-            self.sample_ids = cast(List[SampleId], [i for i in range(self.length)])
+            self.sample_ids = [i for i in range(self.length)]
             self.sample_id_type = int
-            self._grouped = False
         elif self.length is None and self.sample_ids is not None:
-            self._grouped = len(self.sample_ids) > 0 and isinstance(self.sample_ids[0], list)
-            if self._grouped:
-                groups = cast(List[List[SampleId]], self.sample_ids)
-                assert len(groups) >= 1, "Grouped PreprocessResponse must have at least one group."
-                for group in groups:
-                    assert isinstance(group, list) and len(group) > 0, \
-                        "Each group in a grouped PreprocessResponse must be a non-empty list."
-                self.sample_id_type = type(groups[0][0])
-                for group in groups:
-                    for sample_id in group:
-                        assert isinstance(sample_id, self.sample_id_type), \
-                            f"Sample id should be of type {self.sample_id_type.__name__}. Got: {type(sample_id)}"
-                self.length = sum(len(group) for group in groups)
-            else:
-                flat_ids = cast(List[SampleId], self.sample_ids)
-                self.length = len(flat_ids)
-                if self.sample_id_type is None:
-                    self.sample_id_type = str
-                if self.sample_id_type == str:
-                    for sample_id in flat_ids:
-                        assert isinstance(sample_id, str), f"Sample id should be of type str. Got: {type(sample_id)}"
+            self.length = len(self.sample_ids)
+            if self.sample_id_type is None:
+                self.sample_id_type = str
+            if self.sample_id_type == str:
+                for sample_id in self.sample_ids:
+                    assert isinstance(sample_id, str), f"Sample id should be of type str. Got: {type(sample_id)}"
         else:
             raise Exception("length is deprecated, please use sample_ids instead.")
 
@@ -98,36 +79,9 @@ class PreprocessResponse:
         return id(self)
 
     def __len__(self) -> int:
-        assert self.length is not None
-        return self.length
-
-    @property
-    def is_grouped(self) -> bool:
-        return self._grouped
-
-    @property
-    def num_groups(self) -> int:
-        # Grouped: the number of groups. Flat: the (flat) sample count, since sample_ids is the
-        # flat id list and there is one sample per "group".
         assert self.sample_ids is not None
         return len(self.sample_ids)
 
-    @property
-    def groups(self) -> Optional[List[List[SampleId]]]:
-        # The nested list of groups when grouped, otherwise None. Callers should
-        # branch on is_grouped before relying on this.
-        return cast(List[List[SampleId]], self.sample_ids) if self._grouped else None
-
-    @property
-    def flat_sample_ids(self) -> List[SampleId]:
-        # All sample ids in order: concatenation of the groups when grouped,
-        # otherwise the flat sample_ids list as-is.
-        assert self.sample_ids is not None
-        if self._grouped:
-            groups = cast(List[List[SampleId]], self.sample_ids)
-            return [sample_id for group in groups for sample_id in group]
-        return cast(List[SampleId], self.sample_ids)
-
 @dataclass
 class ElementInstance:
     name: str
@@ -138,6 +92,16 @@ SectionCallableInterface = Callable[[Union[int, str], PreprocessResponse], npt.N
 InstanceCallableInterface = Callable[[Union[int, str], PreprocessResponse, int], Optional[ElementInstance]]
 InstanceLengthCallableInterface = Callable[[Union[int, str], PreprocessResponse], int]
 
+# (sample_id, prev_inputs, prev_outputs, preprocess) -> next model inputs, or None to end the chain.
+# First call per chain receives prev_inputs=None, prev_outputs=None and returns the initial inputs.
+# Keys starting with '_' are passthrough state: never fed to the model, handed back verbatim in
+# prev_inputs on the next call. The hook must be stateless and deterministically seeded (the engine
+# may replay a step after crash recovery and relies on identical results).
+AutoregressiveStepCallableInterface = Callable[
+    [Union[int, str], Optional[Dict[str, npt.NDArray[np.float32]]], Optional[Dict[str, npt.NDArray[np.float32]]],
+     PreprocessResponse],
+    Optional[Dict[str, npt.NDArray[np.float32]]]]
+
 MetadataSectionCallableInterface = Union[
     Callable[[Union[int, str], PreprocessResponse], int],
     Callable[[Union[int, str], PreprocessResponse], Dict[str, int]],
@@ -285,6 +249,23 @@ class CustomLatentSpaceHandler:
     function: SectionCallableInterface
     name: str = 'custom_latent_space'
 
+
+# How a chain's latent-space vectors are derived from its per-step forward passes.
+# 'last_step': every latent space comes from the final step's forward pass, except input-kind
+# latent spaces which come from the first step (the original sample, before generated content
+# dominates the model inputs). 'mean': every latent space is the elementwise mean over all steps.
+AUTOREGRESSIVE_LATENT_SPACE_AGGREGATIONS = ('last_step', 'mean')
+
+
+@dataclass
+class AutoregressiveStepHandler:
+    function: AutoregressiveStepCallableInterface
+    name: str = 'autoregressive_step'
+    # Per-model-input shapes (unbatched, underscore passthrough keys excluded), discovered by running
+    # the hook's first call at parse time. Fills the role InputHandler.shape plays for input encoders.
+    input_shapes: Optional[Dict[str, List[int]]] = None
+    latent_space_aggregation: str = 'last_step'
+
 @dataclass
 class PredictionTypeHandler:
     name: str
@@ -325,6 +306,7 @@ class DatasetIntegrationSetup:
     custom_layers: Dict[str, CustomLayerHandler] = field(default_factory=dict)
     custom_latent_space: Optional[CustomLatentSpaceHandler] = None
     simulations: List[SimulationHandler] = field(default_factory=list)
+    autoregressive_step: Optional[AutoregressiveStepHandler] = None
 
 
 @dataclass

code_loader/inner_leap_binder/leapbinder.py

@@ -13,7 +13,8 @@ from code_loader.contract.datasetclasses import SectionCallableInterface, InputH
     CustomMultipleReturnCallableInterfaceMultiArgs, DatasetBaseHandler, custom_latent_space_attribute, \
     RawInputsForHeatmap, VisualizerHandlerData, MetricHandlerData, CustomLossHandlerData, SamplePreprocessResponse, \
     ElementInstanceMasksHandler, InstanceCallableInterface, CustomLatentSpaceHandler, InstanceMetricHandler, \
-    SimulationHandler, _simulation_context
+    SimulationHandler, _simulation_context, AutoregressiveStepHandler, AutoregressiveStepCallableInterface, \
+    AUTOREGRESSIVE_LATENT_SPACE_AGGREGATIONS
 from code_loader.contract.enums import LeapDataType, DataStateEnum, DataStateType, MetricDirection, DatasetMetadataType
 from code_loader.contract.mapping import NodeConnection, NodeMapping, NodeMappingType
 from code_loader.contract.responsedataclasses import DatasetTestResultPayload, LeapAnalysisConfiguration
@@ -524,6 +525,34 @@ class LeapBinder:
         """
         self.setup_container.custom_latent_space = CustomLatentSpaceHandler(function)
 
+    def set_autoregressive_step(self, function: AutoregressiveStepCallableInterface,
+                                latent_space_aggregation: str = 'last_step') -> None:
+        """
+        Set the autoregressive step hook — the feedback function that drives a chain:
+        it supplies the model's initial inputs on its first call (prev_inputs=None, prev_outputs=None)
+        and turns each step's inputs/outputs into the next step's inputs until it returns None.
+        An autoregressive integration has no input encoders; the hook is the sole input source.
+        latent_space_aggregation declares how the chain's latent-space vectors are derived from
+        its steps (see AUTOREGRESSIVE_LATENT_SPACE_AGGREGATIONS).
+        """
+        if self.setup_container.autoregressive_step is not None:
+            raise Exception('tensorleap_autoregressive_step is already defined. '
+                            'Only one autoregressive step hook is allowed per integration.')
+        if latent_space_aggregation not in AUTOREGRESSIVE_LATENT_SPACE_AGGREGATIONS:
+            raise Exception(f'tensorleap_autoregressive_step: unknown latent_space_aggregation '
+                            f'{latent_space_aggregation!r}. Supported values: '
+                            f'{", ".join(AUTOREGRESSIVE_LATENT_SPACE_AGGREGATIONS)}.')
+        self.setup_container.autoregressive_step = AutoregressiveStepHandler(
+            function, latent_space_aggregation=latent_space_aggregation)
+
+        # Builtin chain metadata, declared at parse time so it survives the reporter's
+        # metadata type mapping; the placeholder values are overwritten by the engine when a
+        # chain finalizes (realized length, truncated-by-safety-cap flag).
+        def builtin_chain_metadata(idx: Any, preprocess: PreprocessResponse) -> Dict[str, Any]:
+            return {'length': -1, 'truncated': False}
+
+        self.set_metadata(builtin_chain_metadata, 'builtin_chain')
+
     def set_custom_layer(self, custom_layer: Type[Any], name: str, inspect_layer: bool = False,
                          kernel_index: Optional[int] = None, use_custom_latent_space: bool = False) -> None:
         """
@@ -606,15 +635,6 @@ class LeapBinder:
             if state_enum in preprocess_result_dict:
                 preprocess_result_dict_in_correct_order[state_enum] = preprocess_result_dict[state_enum]
 
-        # All-or-none: every state must agree on grouped (nested sample_ids) vs flat.
-        if len({r.is_grouped for r in preprocess_result_dict_in_correct_order.values()}) > 1:
-            shapes = ', '.join(
-                f"{state.name}={'grouped' if r.is_grouped else 'flat'}"
-                for state, r in preprocess_result_dict_in_correct_order.items())
-            raise Exception(
-                "All preprocess states must be either all grouped (nested sample_ids) or all flat, "
-                f"but got a mix: {shapes}.")
-
         return preprocess_result_dict_in_correct_order
 
     def get_preprocess_unlabeled_result(self) -> Optional[PreprocessResponse]:
@@ -635,13 +655,7 @@ class LeapBinder:
             preprocess_response: PreprocessResponse, test_result: List[DatasetTestResultPayload],
             dataset_base_handler: Union[DatasetBaseHandler, MetadataHandler], state: DataStateEnum) -> List[DatasetTestResultPayload]:
         assert preprocess_response.sample_ids is not None
-        if preprocess_response.is_grouped:
-            # Grouped: probe with the first group, then reduce to a single sample's result so the
-            # recorded shape/type is per-sample (matching the flat case), not the (B, *) batch.
-            group = preprocess_response.sample_ids[0]
-            raw_result = dataset_base_handler.function(group, preprocess_response)[0]
-        else:
-            raw_result = dataset_base_handler.function(preprocess_response.sample_ids[0], preprocess_response)
+        raw_result = dataset_base_handler.function(preprocess_response.sample_ids[0], preprocess_response)
         handler_type = 'metadata' if isinstance(dataset_base_handler, MetadataHandler) else None
         if isinstance(dataset_base_handler, MetadataHandler):
             if isinstance(raw_result, dict):
@@ -749,6 +763,7 @@ class LeapBinder:
                     )
 
     def check(self) -> None:
+        self.validate_autoregressive_setup()
         preprocess_result = self.get_preprocess_result()
         self.check_preprocess(preprocess_result)
         self.check_handlers(preprocess_result)
@@ -756,6 +771,43 @@ class LeapBinder:
         self.validate_ignore_latent_spaces()
         print("Successful!")
 
+    def validate_autoregressive_setup(self) -> None:
+        """Structural rules for autoregressive integrations (no-op when the hook is absent).
+
+        Order-independent: runs after the whole script registered, because decorators may
+        appear in any order in the integration file.
+        """
+        if self.setup_container.autoregressive_step is None:
+            return
+        if self.setup_container.inputs:
+            raise Exception(
+                "An autoregressive integration must not declare input encoders — "
+                "tensorleap_autoregressive_step is the sole source of model inputs "
+                "(its first call returns the initial inputs). Found input encoders: "
+                f"{[handler.name for handler in self.setup_container.inputs]}.")
+        if self.setup_container.simulations:
+            raise Exception(
+                "Simulations are not supported together with tensorleap_autoregressive_step yet: "
+                "the autoregressive entry points do not resolve synthetic sample ids. Remove the "
+                "simulation decorators or the autoregressive hook.")
+        if self.setup_container.instance_masks:
+            raise Exception(
+                "Element instances are not supported together with tensorleap_autoregressive_step: "
+                "instance generation masks the sample's encoded inputs, which do not exist in an "
+                "autoregressive integration. Remove the instance encoders or the autoregressive hook.")
+        if self.setup_container.custom_latent_space is not None:
+            raise Exception(
+                "tensorleap_custom_latent_space is not supported together with "
+                "tensorleap_autoregressive_step: the custom latent space function only sees "
+                "(sample_id, preprocess) — it runs before generation and cannot observe the chain. "
+                "The chain's latent is the final step's model latent.")
+        if not self.setup_container.prediction_types:
+            raise Exception(
+                "An autoregressive integration must declare prediction types "
+                "(tensorleap_load_model(prediction_types=[...])): the engine keys each step's "
+                "model outputs by the declared names to build the hook's prev_outputs dict. "
+                "Without them, prev_outputs would be empty at runtime.")
+
     def validate_ignore_latent_spaces(self) -> None:
         """Validate leap_analysis_configuration.ignore_latent_spaces against this binder.