code-loader 1.0.197.dev0__tar.gz → 1.0.197.dev2__tar.gz

{code_loader-1.0.197.dev0 → code_loader-1.0.197.dev2}/PKG-INFO

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

{code_loader-1.0.197.dev0 → code_loader-1.0.197.dev2}/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
+from typing import Any, Callable, List, Optional, Dict, Union, Type, Literal, Tuple, cast
 import re
 import numpy as np
 import numpy.typing as npt
@@ -15,6 +15,8 @@ custom_latent_space_attribute = "custom_latent_space"
 
 _simulation_context: Dict[str, bool] = {"active": False}
 
+SampleId = Union[int, str]
+
 
 @dataclass
 class PreprocessResponse:
@@ -39,27 +41,54 @@ class PreprocessResponse:
     """
     length: Optional[int] = None  # Deprecated. Please use sample_ids instead
     data: Any = None
-    sample_ids: Optional[Union[List[str], List[int]]] = None
+    sample_ids: Optional[Union[List[SampleId], List[List[SampleId]]]] = 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 = [i for i in range(self.length)]
+            self.sample_ids = cast(List[SampleId], [i for i in range(self.length)])
             self.sample_id_type = int
-        elif self.length is None and self.sample_ids is not None:
-            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)}"
+            self._grouped = False
+        elif self.sample_ids is not None:
+            given_length = self.length
+            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])
+                assert issubclass(self.sample_id_type, (str, int)), \
+                    f"Sample id should be of type str or int. Got: {self.sample_id_type.__name__}"
+                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)}"
+            # dataclasses.replace() re-invokes __init__ with both fields already populated
+            # (length previously derived from sample_ids); only reject a genuine mismatch.
+            if given_length is not None:
+                assert given_length == self.length, \
+                    f"Inconsistent PreprocessResponse: length={given_length} but sample_ids implies {self.length}."
         else:
             raise Exception("length is deprecated, please use sample_ids instead.")
 
@@ -79,9 +108,36 @@ 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
@@ -252,6 +308,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.

{code_loader-1.0.197.dev0 → code_loader-1.0.197.dev2}/code_loader/inner_leap_binder/leapbinder.py

@@ -1,6 +1,9 @@
 
+import builtins
 import inspect
-from typing import Callable, List, Optional, Dict, Any, Type, Union, get_args
+import os
+from contextlib import contextmanager
+from typing import Callable, List, Optional, Dict, Any, Type, Union, get_args, cast, Iterator, Set
 
 import numpy as np
 import numpy.typing as npt
@@ -16,7 +19,8 @@ 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
+from code_loader.contract.enums import LeapDataType, DataStateEnum, DataStateType, MetricDirection, DatasetMetadataType, \
+    TestingSectionEnum
 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
@@ -33,6 +37,25 @@ 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
@@ -514,7 +537,8 @@ class LeapBinder:
         self.setup_container.metadata.append(MetadataHandler(name, function, metadata_type))
 
     def set_custom_latent_space(self, function: SectionCallableInterface,
-                                name: Optional[str] = None) -> None:
+                                name: Optional[str] = None,
+                                use_ls_for_analysis: bool = False) -> None:
         """
         Register a custom latent space function.
 
@@ -531,6 +555,10 @@ class LeapBinder:
             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.
         """
         if name is None:
             name = custom_latent_space_attribute
@@ -540,7 +568,22 @@ class LeapBinder:
                 f"@tensorleap_custom_latent_space must have a unique name "
                 f"(pass name='...' to distinguish them)."
             )
-        self.setup_container.custom_latent_spaces[name] = CustomLatentSpaceHandler(function, name)
+        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:
@@ -565,7 +608,9 @@ 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) -> Dict[str, Any]:
+        def builtin_chain_metadata(idx: Any, preprocess: PreprocessResponse) -> Any:
+            if isinstance(idx, list):
+                return [{'length': -1, 'truncated': False} for _ in idx]
             return {'length': -1, 'truncated': False}
 
         self.set_metadata(builtin_chain_metadata, 'builtin_chain')
@@ -729,6 +774,15 @@ 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]:
@@ -749,7 +803,34 @@ 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
-        raw_result = dataset_base_handler.function(preprocess_response.sample_ids[0], preprocess_response)
+        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
+
         handler_type = 'metadata' if isinstance(dataset_base_handler, MetadataHandler) else None
         if isinstance(dataset_base_handler, MetadataHandler):
             if isinstance(raw_result, dict):
@@ -789,6 +870,7 @@ 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:
@@ -810,6 +892,7 @@ 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:
@@ -845,10 +928,15 @@ 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 = [0]
+            if preprocess_response.is_grouped:
+                raise Exception(
+                    "Simulation '{}' returned a grouped PreprocessResponse; simulations must "
+                    "return a flat (non-grouped) PreprocessResponse.".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)
             for handler in self.setup_container.inputs:
-                out1 = handler.function(preprocess_response.sample_ids[0], preprocess_response)
-                out2 = handler.function(preprocess_response.sample_ids[0], preprocess_response)
+                out1 = handler.function(sim_sample_ids[0], preprocess_response)
+                out2 = handler.function(sim_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(
@@ -906,7 +994,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 "