code-loader 1.0.196.dev1__tar.gz → 1.0.197.dev1__tar.gz

{code_loader-1.0.196.dev1 → code_loader-1.0.197.dev1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: code-loader
-Version: 1.0.196.dev1
+Version: 1.0.197.dev1
 Summary: 
 Home-page: https://github.com/tensorleap/code-loader
 License: MIT

{code_loader-1.0.196.dev1 → code_loader-1.0.197.dev1}/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, Tuple, cast
+from typing import Any, Callable, List, Optional, Dict, Union, Type, Literal, Tuple
 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,46 +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)
-    _group_pos_cache: Optional[Dict[SampleId, Tuple[List[SampleId], int]]] = field(
-        default=None, 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.")
 
@@ -100,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
@@ -300,6 +252,7 @@ class MetadataHandler:
 class CustomLatentSpaceHandler:
     function: SectionCallableInterface
     name: str = 'custom_latent_space'
+    use_ls_for_analysis: bool = False
 
 
 # How a chain's latent-space vectors are derived from its per-step forward passes.
@@ -382,7 +335,7 @@ class DatasetIntegrationSetup:
     metrics: List[MetricHandler] = field(default_factory=list)
     instance_metrics: List[InstanceMetricHandler] = field(default_factory=list)
     custom_layers: Dict[str, CustomLayerHandler] = field(default_factory=dict)
-    custom_latent_space: Optional[CustomLatentSpaceHandler] = None
+    custom_latent_spaces: Dict[str, CustomLatentSpaceHandler] = field(default_factory=dict)
     simulations: List[SimulationHandler] = field(default_factory=list)
     autoregressive_step: Optional[AutoregressiveStepHandler] = None
     autoregressive_metrics: List[AutoregressiveMetricHandler] = field(default_factory=list)
@@ -398,6 +351,6 @@ class DatasetSample:
     metadata_is_none: Dict[str, bool]
     index: Union[int, str]
     state: DataStateEnum
-    custom_latent_space: Optional[npt.NDArray[np.float32]] = None
+    custom_latent_spaces: Optional[Dict[str, npt.NDArray[np.float32]]] = None
     instance_masks: Optional[Dict[str, ElementInstance]] = None
 

{code_loader-1.0.196.dev1 → code_loader-1.0.197.dev1}/code_loader/inner_leap_binder/leapbinder.py

@@ -1,9 +1,6 @@
 
-import builtins
 import inspect
-import os
-from contextlib import contextmanager
-from typing import Callable, List, Optional, Dict, Any, Type, Union, get_args, cast, Iterator, Set
+from typing import Callable, List, Optional, Dict, Any, Type, Union, get_args
 
 import numpy as np
 import numpy.typing as npt
@@ -19,8 +16,7 @@ from code_loader.contract.datasetclasses import SectionCallableInterface, InputH
     SimulationHandler, _simulation_context, AutoregressiveStepHandler, AutoregressiveStepCallableInterface, \
     AUTOREGRESSIVE_LATENT_SPACE_AGGREGATIONS, AUTOREGRESSIVE_IMPLICIT_ARG_NAMES, \
     AutoregressiveMetricHandler, AutoregressiveLossHandler, AutoregressiveVisualizerHandler
-from code_loader.contract.enums import LeapDataType, DataStateEnum, DataStateType, MetricDirection, DatasetMetadataType, \
-    TestingSectionEnum
+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
 from code_loader.contract.visualizer_classes import map_leap_data_type_to_visualizer_class
@@ -37,25 +33,6 @@ from code_loader.visualizers.default_visualizers import DefaultVisualizer, \
 mapping_runtime_mode_env_var_mame = '__MAPPING_RUNTIME_MODE__'
 
 
-@contextmanager
-def _track_opened_files() -> Iterator[Set[str]]:
-    opened: Set[str] = set()
-    original_open = builtins.open
-
-    def tracking_open(file: Any, *args: Any, **kwargs: Any) -> Any:
-        mode = args[0] if args else kwargs.get('mode', 'r')
-        is_read_mode = isinstance(mode, str) and not any(c in mode for c in ('a', 'w', 'x'))
-        if is_read_mode and isinstance(file, (str, os.PathLike)):
-            opened.add(os.fspath(file))
-        return original_open(file, *args, **kwargs)
-
-    builtins.open = tracking_open
-    try:
-        yield opened
-    finally:
-        builtins.open = original_open
-
-
 def _stringized_annotation_type_name(annotation: Any) -> Optional[str]:
     """Bare type name if ``annotation`` is a stringized annotation (from ``from __future__
     import annotations`` or a quoted hint), else ``None``. code_loader inspects raw
@@ -536,18 +513,54 @@ class LeapBinder:
         """
         self.setup_container.metadata.append(MetadataHandler(name, function, metadata_type))
 
-    def set_custom_latent_space(self, function: SectionCallableInterface) -> None:
+    def set_custom_latent_space(self, function: SectionCallableInterface,
+                                name: Optional[str] = None,
+                                use_ls_for_analysis: bool = False) -> None:
         """
-        Set a custom latent space function.
+        Register a custom latent space function.
+
+        Multiple custom latent spaces may be registered as long as each has a
+        unique name. They are stored name-keyed (insertion order preserved), so a
+        later registration no longer overrides an earlier one.
 
         Args:
-        function (SectionCallableInterface): The metadata handler function.
+        function (SectionCallableInterface): The latent-space handler function.
             This function receives:
                 subset (PreprocessResponse): The subset of the data.
                 index (int): The index of the sample within the subset.
-            This function should numpy float32 array contains the latent space vec of the sample.
+            This function should return a numpy float32 array containing the latent
+            space vec of the sample.
+        name (Optional[str]): Unique name for this custom latent space. Defaults to
+            the reserved single-LS name for backward compatibility.
+        use_ls_for_analysis (bool): When True, the engine uses this custom latent
+            space for the Out-Of-Distribution and Domain-Gap insights instead of the
+            built-in defaults. At most one registered custom latent space may set this;
+            registering a second one with the flag raises.
         """
-        self.setup_container.custom_latent_space = CustomLatentSpaceHandler(function)
+        if name is None:
+            name = custom_latent_space_attribute
+        if name in self.setup_container.custom_latent_spaces:
+            raise Exception(
+                f"A custom latent space named '{name}' is already registered. Each "
+                f"@tensorleap_custom_latent_space must have a unique name "
+                f"(pass name='...' to distinguish them)."
+            )
+        if use_ls_for_analysis:
+            already_flagged = [
+                existing_name
+                for existing_name, handler in self.setup_container.custom_latent_spaces.items()
+                if handler.use_ls_for_analysis
+            ]
+            if already_flagged:
+                raise Exception(
+                    f"use_ls_for_analysis=True is already set on custom latent space "
+                    f"'{already_flagged[0]}'. Only one custom latent space may set "
+                    f"use_ls_for_analysis=True (it selects the latent space used for the "
+                    f"Out-Of-Distribution and Domain-Gap insights). Set it on '{name}' "
+                    f"or '{already_flagged[0]}', not both."
+                )
+        self.setup_container.custom_latent_spaces[name] = CustomLatentSpaceHandler(
+            function, name, use_ls_for_analysis)
 
     def set_autoregressive_step(self, function: AutoregressiveStepCallableInterface,
                                 latent_space_aggregation: str = 'last_step') -> None:
@@ -572,9 +585,7 @@ class LeapBinder:
         # 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) -> Any:
-            if isinstance(idx, list):
-                return [{'length': -1, 'truncated': False} for _ in idx]
+        def builtin_chain_metadata(idx: Any, preprocess: PreprocessResponse) -> Dict[str, Any]:
             return {'length': -1, 'truncated': False}
 
         self.set_metadata(builtin_chain_metadata, 'builtin_chain')
@@ -738,15 +749,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]:
@@ -767,34 +769,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
-        opened_files: Optional[Set[str]] = 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.
-            # (casts: sample_ids[0] is a group list here, and the grouped result is a per-sample list.)
-            group = preprocess_response.sample_ids[0]
-            with _track_opened_files() as opened_files:
-                raw_result = cast(Any, dataset_base_handler.function(cast(Any, group), preprocess_response))[0]
-        else:
-            raw_result = dataset_base_handler.function(
-                cast(Union[int, str], preprocess_response.sample_ids[0]), preprocess_response)
-
-        def _warn_if_group_spans_multiple_files(payloads: List[DatasetTestResultPayload]) -> None:
-            if opened_files is None or len(opened_files) <= 1:
-                return
-            shown = sorted(opened_files)[:5]
-            suffix = f" (+{len(opened_files) - 5} more)" if len(opened_files) > 5 else ""
-            warning = (
-                f"Encoding declared group 0 (size {len(cast(List[Any], group))}) for '{dataset_base_handler.name}' opened "
-                f"{len(opened_files)} distinct files: {shown}{suffix}. If this group is meant to be one "
-                f"physical source (e.g. one Parquet file), it may be grouped incorrectly - check your "
-                f"preprocess grouping logic. (Only files opened via Python's open() are tracked, so this "
-                f"can miss other I/O paths or be a false positive if grouping isn't meant to reflect file "
-                f"locality.)")
-            for payload in payloads:
-                existing = payload.display.get(TestingSectionEnum.Errors.name, '')
-                payload.display[TestingSectionEnum.Errors.name] = (existing + '\n' if existing else '') + warning
-
+        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):
@@ -834,7 +809,6 @@ class LeapBinder:
             else:
                 if raw_result is None:
                     if state != DataStateEnum.training:
-                        _warn_if_group_spans_multiple_files(test_result)
                         return test_result
 
                     if dataset_base_handler.metadata_type is None:
@@ -856,7 +830,6 @@ class LeapBinder:
             # setting shape in setup for all encoders
             if isinstance(dataset_base_handler, (InputHandler, GroundTruthHandler)):
                 dataset_base_handler.shape = result_shape
-        _warn_if_group_spans_multiple_files(test_result)
         return test_result
 
     def check_handlers(self, preprocess_result: Dict[DataStateEnum, PreprocessResponse]) -> None:
@@ -892,11 +865,10 @@ class LeapBinder:
             preprocess_response.tl_generated = True
             if not preprocess_response.length or preprocess_response.length < 1:
                 raise Exception("Simulation '{}' returned PreprocessResponse with length < 1".format(sim.name))
-            preprocess_response.sample_ids = cast(List[Union[int, str]], [0])
-            sim_sample_ids = cast(List[Union[int, str]], preprocess_response.sample_ids)
+            preprocess_response.sample_ids = [0]
             for handler in self.setup_container.inputs:
-                out1 = handler.function(sim_sample_ids[0], preprocess_response)
-                out2 = handler.function(sim_sample_ids[0], preprocess_response)
+                out1 = handler.function(preprocess_response.sample_ids[0], preprocess_response)
+                out2 = handler.function(preprocess_response.sample_ids[0], preprocess_response)
                 if not np.allclose(out1, out2):
                     raise Exception(
                         "Simulation '{}': encoder '{}' is non-deterministic — consecutive calls with seed=0 returned different outputs".format(
@@ -954,7 +926,7 @@ class LeapBinder:
                 "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:
+        if self.setup_container.custom_latent_spaces:
             raise Exception(
                 "tensorleap_custom_latent_space is not supported together with "
                 "tensorleap_autoregressive_step: the custom latent space function only sees "