code-loader 1.0.199__tar.gz → 1.0.201__tar.gz

{code_loader-1.0.199 → code_loader-1.0.201}/PKG-INFO

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

{code_loader-1.0.199 → code_loader-1.0.201}/code_loader/contract/datasetclasses.py

@@ -8,7 +8,7 @@ import numpy.typing as npt
 from code_loader.contract.enums import DataStateType, DataStateEnum, LeapDataType, ConfusionMatrixValue, \
     MetricDirection, DatasetMetadataType
 from code_loader.contract.visualizer_classes import LeapImage, LeapText, LeapGraph, LeapHorizontalBar, \
-    LeapTextMask, LeapImageMask, LeapImageWithBBox, LeapImageWithHeatmap, LeapVideo
+    LeapTextMask, LeapImageMask, LeapImageWithBBox, LeapImageWithHeatmap, LeapVideo, LeapAudio
 from code_loader.contract.sim_config import SimConfig
 
 custom_latent_space_attribute = "custom_latent_space"
@@ -194,11 +194,12 @@ VisualizerCallableInterface = Union[
     Callable[..., LeapImageMask],
     Callable[..., LeapTextMask],
     Callable[..., LeapImageWithBBox],
-    Callable[..., LeapImageWithHeatmap]
+    Callable[..., LeapImageWithHeatmap],
+    Callable[..., LeapAudio]
 ]
 
 LeapData = Union[LeapImage, LeapText, LeapGraph, LeapHorizontalBar, LeapImageMask, LeapTextMask, LeapImageWithBBox,
-LeapImageWithHeatmap, LeapVideo]
+LeapImageWithHeatmap, LeapVideo, LeapAudio]
 
 CustomCallableInterface = Callable[..., Any]
 
@@ -308,6 +309,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.
@@ -390,7 +392,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)
@@ -406,6 +408,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.199 → code_loader-1.0.201}/code_loader/contract/visualizer_classes.py

@@ -357,16 +357,14 @@ class LeapAudio:
         should be mixed down before constructing LeapAudio.
     sample_rate (int): Samples per second — needed to encode the clip and to map
         the playhead position to time.
-    visual (Union[LeapImage, LeapGraph]): The thing you look at.
-    x_range (Optional[Tuple[float, float]]): Time extent (in seconds) that the
-        visual's horizontal axis spans, used to align the tick. Defaults to
-        (0, N / sample_rate).
+    visual (Union[LeapImage, LeapGraph]): The thing you look at. For a waveform
+        the nested LeapGraph carries its own x_range (time span); the playhead is
+        driven by the audio element's currentTime, so no time range is needed here.
     type (LeapDataType): The data type, default is LeapDataType.Audio.
     """
     audio: npt.NDArray
     sample_rate: int
     visual: Union[LeapImage, LeapGraph]
-    x_range: Optional[Tuple[float, float]] = None
     type: LeapDataType = LeapDataType.Audio
 
     def __post_init__(self) -> None:
@@ -377,11 +375,6 @@ class LeapAudio:
                       'Audio visual must be a LeapImage (spectrogram) or LeapGraph (waveform)')
         if not isinstance(self.sample_rate, int) or self.sample_rate <= 0:
             raise LeapValidationError(f'sample_rate must be a positive int, got {self.sample_rate}')
-        if self.x_range is not None:
-            if len(self.x_range) != 2:
-                raise LeapValidationError('x_range must be a tuple of length 2')
-            validate_type(type(self.x_range[0]), [float, int], 'x_range must be a tuple of floats or integers')
-            validate_type(type(self.x_range[1]), [float, int], 'x_range must be a tuple of floats or integers')
 
 
 map_leap_data_type_to_visualizer_class = {

{code_loader-1.0.199 → code_loader-1.0.201}/code_loader/inner_leap_binder/leapbinder.py

@@ -536,18 +536,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:
@@ -958,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 "

{code_loader-1.0.199 → code_loader-1.0.201}/code_loader/inner_leap_binder/leapbinder_decorators.py

@@ -1703,8 +1703,14 @@ def tensorleap_metadata(
     return decorating_function
 
 
-def tensorleap_custom_latent_space():
+def tensorleap_custom_latent_space(name: Optional[str] = None, use_ls_for_analysis: bool = False):
+    assert isinstance(use_ls_for_analysis, bool), \
+        ("tensorleap_custom_latent_space validation failed: use_ls_for_analysis must be a bool. "
+         f"Got {type(use_ls_for_analysis)}.")
+
     def decorating_function(user_function: SectionCallableInterface):
+        ls_name = name if name is not None else user_function.__name__
+
         def _validate_input_args(sample_id: Union[int, str, list], preprocess_response: PreprocessResponse):
             _validate_id_or_group(sample_id, preprocess_response, 'tensorleap_custom_latent_space')
 
@@ -1715,9 +1721,9 @@ def tensorleap_custom_latent_space():
             if single_result.ndim > 1:
                 flat_dim = int(np.prod(single_result.shape))
                 store_general_warning(
-                    key=("tensorleap_custom_latent_space_flatten", tuple(single_result.shape)),
+                    key=("tensorleap_custom_latent_space_flatten", ls_name, tuple(single_result.shape)),
                     message=(
-                        f"tensorleap_custom_latent_space returned per-sample shape {tuple(single_result.shape)} "
+                        f"tensorleap_custom_latent_space '{ls_name}' returned per-sample shape {tuple(single_result.shape)} "
                         f"(ndim={single_result.ndim}). Tensorleap assumes per-sample shape (d, ...) and will "
                         f"flatten to ({flat_dim},) before downstream visualization and clustering. "
                         f"If you want a different aggregation (e.g. global average pooling), do it "
@@ -1760,7 +1766,8 @@ def tensorleap_custom_latent_space():
 
             return result
 
-        leap_binder.set_custom_latent_space(inner_without_validate)
+        leap_binder.set_custom_latent_space(inner_without_validate, ls_name,
+                                            use_ls_for_analysis=use_ls_for_analysis)
 
         def inner(sample_id, preprocess_response):
             if os.environ.get(mapping_runtime_mode_env_var_mame):

{code_loader-1.0.199 → code_loader-1.0.201}/code_loader/leaploader.py

@@ -283,17 +283,7 @@ class LeapLoader(LeapLoaderBase):
 
         metadata, metadata_is_none = self.get_metadata(state, sample_id)
 
-        custom_latent_space = None
-        if global_leap_binder.setup_container.custom_latent_space is not None:
-            latent_fn = global_leap_binder.setup_container.custom_latent_space.function
-            preprocess_state = preprocess_result[state]
-            if preprocess_state.is_grouped:
-                # Single-sample fetch: encode this sample only, not the whole group (same
-                # memory rationale as _get_dataset_handlers; see grouped-fetch-oom-bug.md).
-                self._locate_group(preprocess_state, sample_id)  # validates group membership
-                custom_latent_space = self._to_grouped_list(latent_fn([sample_id], preprocess_state))[0]
-            else:
-                custom_latent_space = latent_fn(sample_id, preprocess_state)
+        custom_latent_spaces = self._get_custom_latent_spaces(sample_id, preprocess_result[state])
         instance_mask = self._get_instances_masks(state, sample_id, instance_id)
         sample = DatasetSample(inputs=self._get_inputs(state, sample_id),
                                gt=None if state == DataStateEnum.unlabeled else self._get_gt(state, sample_id),
@@ -301,7 +291,7 @@ class LeapLoader(LeapLoaderBase):
                                metadata_is_none=metadata_is_none,
                                index=sample_id,
                                state=state,
-                               custom_latent_space=custom_latent_space,
+                               custom_latent_spaces=custom_latent_spaces,
                                instance_masks=instance_mask)
         return sample
 
@@ -332,11 +322,7 @@ class LeapLoader(LeapLoaderBase):
                     handler.name, handler_result
                 )
 
-        custom_latent_space = None
-        if global_leap_binder.setup_container.custom_latent_space is not None:
-            custom_latent_space = global_leap_binder.setup_container.custom_latent_space.function(
-                original_sample_id, preprocess
-            )
+        custom_latent_spaces = self._get_custom_latent_spaces(original_sample_id, preprocess)
 
         return DatasetSample(
             inputs=inputs,
@@ -345,7 +331,7 @@ class LeapLoader(LeapLoaderBase):
             metadata_is_none=metadata_is_none,
             index=synthetic_index,
             state=DataStateEnum.additional,
-            custom_latent_space=custom_latent_space,
+            custom_latent_spaces=custom_latent_spaces,
             instance_masks=None,
         )
 
@@ -1114,17 +1100,20 @@ class LeapLoader(LeapLoaderBase):
                 metadata[handler.name] = [v for v, _ in converted]
                 metadata_is_none[handler.name] = [is_none for _, is_none in converted]
 
-        # custom_latent_space is group-aware like the input/GT encoders: hand it the whole group in a
-        # single call so a file-backed latent fn loads the file once, and normalize to (B, d). This
-        # keeps the mandatory group path non-lossy. instance_masks stay None here: they are a
-        # per-(sample, instance) concern that needs an instance_id the group fetch does not carry.
-        custom_latent_space = None
-        if global_leap_binder.setup_container.custom_latent_space is not None:
-            latent_fn = global_leap_binder.setup_container.custom_latent_space.function
-            custom_latent_space = self._to_grouped_list(latent_fn(group_ids, preprocess_state))
+        # custom_latent_spaces are group-aware like the input/GT encoders: hand each the whole group
+        # in a single call so a file-backed latent fn loads the file once, and normalize to per-sample
+        # lists. This keeps the mandatory group path non-lossy. instance_masks stay None here: they are
+        # a per-(sample, instance) concern that needs an instance_id the group fetch does not carry.
+        latent_handlers = global_leap_binder.setup_container.custom_latent_spaces
+        custom_latent_spaces = None
+        if latent_handlers:
+            custom_latent_spaces = {
+                name: self._to_grouped_list(handler.function(group_ids, preprocess_state))
+                for name, handler in latent_handlers.items()
+            }
 
         return DatasetSample(inputs=inputs, gt=gt, metadata=metadata, metadata_is_none=metadata_is_none,
-                             index=list(group_ids), state=state, custom_latent_space=custom_latent_space,
+                             index=list(group_ids), state=state, custom_latent_spaces=custom_latent_spaces,
                              instance_masks=None)
 
     def _get_inputs(self, state: DataStateEnum, sample_id: Union[int, str]) -> Dict[str, npt.NDArray[np.float32]]:
@@ -1337,10 +1326,48 @@ class LeapLoader(LeapLoaderBase):
 
         return id_type
 
+    def _get_custom_latent_spaces(
+            self,
+            sample_id: Union[int, str],
+            preprocess: "PreprocessResponse") -> Optional[Dict[str, npt.NDArray[np.float32]]]:
+        handlers = global_leap_binder.setup_container.custom_latent_spaces
+        if not handlers:
+            return None
+        if preprocess.is_grouped:
+            # Single-sample fetch: encode this sample only, not the whole group (same
+            # memory rationale as _get_dataset_handlers; see grouped-fetch-oom-bug.md).
+            self._locate_group(preprocess, sample_id)  # validates group membership
+            return {name: self._to_grouped_list(handler.function([sample_id], preprocess))[0]
+                    for name, handler in handlers.items()}
+        return {name: handler.function(sample_id, preprocess) for name, handler in handlers.items()}
+
     @lru_cache()
     def has_custom_latent_space_decorator(self) -> bool:
         self.exec_script()
-        return global_leap_binder.setup_container.custom_latent_space is not None
+        return len(global_leap_binder.setup_container.custom_latent_spaces) > 0
+
+    @lru_cache()
+    def get_custom_latent_space_names(self) -> Tuple[str, ...]:
+        """Ordered names of all registered custom latent spaces.
+
+        Registration order is the canonical index mapping consumed by the engine
+        (name i -> user_custom_i). Returns a tuple so the lru_cache value is hashable.
+        """
+        self.exec_script()
+        return tuple(global_leap_binder.setup_container.custom_latent_spaces.keys())
+
+    @lru_cache()
+    def get_custom_latent_space_for_analysis(self) -> Optional[str]:
+        """Name of the custom latent space flagged with use_ls_for_analysis=True, if any.
+
+        The engine uses this latent space for the Out-Of-Distribution and Domain-Gap
+        insights instead of the built-in defaults. Returns None when no custom latent
+        space set the flag. At most one is flagged (enforced at registration time)."""
+        self.exec_script()
+        for name, handler in global_leap_binder.setup_container.custom_latent_spaces.items():
+            if handler.use_ls_for_analysis:
+                return name
+        return None
 
     @lru_cache()
     def has_autoregressive_step(self) -> bool:

{code_loader-1.0.199 → code_loader-1.0.201}/code_loader/leaploaderbase.py

@@ -218,6 +218,14 @@ class LeapLoaderBase:
         raise NotImplementedError(f'{type(self).__name__} does not implement '
                                   'autoregressive_visualizer_by_name.')
 
+    @abstractmethod
+    def get_custom_latent_space_names(self) -> Tuple[str, ...]:
+        pass
+
+    @abstractmethod
+    def get_custom_latent_space_for_analysis(self) -> Optional[str]:
+        pass
+
     @abstractmethod
     def get_heatmap_visualizer_raw_vis_input_arg_name(self, visualizer_name: str) -> Optional[str]:
         pass

{code_loader-1.0.199 → code_loader-1.0.201}/code_loader/visualizers/default_visualizers.py

@@ -31,7 +31,6 @@ def default_video_visualizer(data: npt.NDArray[np.float32]) -> LeapVideo:
 
 
 _MAX_WAVEFORM_POINTS = 2000
-_VISUAL_KINDS = ('waveform', 'spectrogram', 'image')
 
 
 def _to_mono(audio: npt.NDArray) -> npt.NDArray:
@@ -54,49 +53,43 @@ def _waveform_graph(series: npt.NDArray, duration: float) -> LeapGraph:
 def default_audio_visualizer(audio: npt.NDArray,
                              sample_rate: int,
                              visual: Optional[npt.NDArray] = None,
-                             visual_kind: Optional[str] = None,
                              time_axis_index: int = -1) -> LeapAudio:
     """Package a playable audio sample plus the visual a playhead runs through.
 
     audio: waveform samples (mono [N], or multi-channel that is mixed down).
     sample_rate: samples per second.
-    visual: the array to display. For 'waveform' a 1-D series; for 'spectrogram'
-        a raw 2-D [freq, time] (or [time, freq]) array to be rendered grayscale;
-        for 'image' an already-rendered [H, W, 1] or [H, W, 3] image (e.g. a
-        spectrogram the caller already ran through their own colormap).
+    visual: the array to display; its kind is deduced from its shape.
+        1-D [N]        -> waveform graph of the samples.
+        2-D [H, W]     -> spectrogram, rendered grayscale.
+        3-D [H, W, C]  -> already-rendered image; C must be 3 (rgb) or 1 (grayscale).
         If omitted, a waveform graph is built from `audio` itself.
-    visual_kind: 'waveform', 'spectrogram', or 'image'. The caller states this
-        explicitly; the shape of `visual` is never inspected to guess it.
-        Required when `visual` is given; ignored when `visual` is omitted.
-    time_axis_index: for 'spectrogram', which axis of `visual` is time (moved to columns).
-        Unused for 'waveform'/'image'.
+    time_axis_index: for a spectrogram, which axis of `visual` is time (moved to columns).
+        Unused otherwise.
     """
     mono = _to_mono(audio)
     duration = mono.shape[0] / float(sample_rate)
-    x_range = (0.0, duration)
 
     if visual is None:  # no visual given: show the raw signal as a waveform
-        return LeapAudio(mono, sample_rate, _waveform_graph(mono, duration), x_range)
-
-    if visual_kind not in _VISUAL_KINDS:
-        raise ValueError(f"visual_kind must be one of {_VISUAL_KINDS}, got {visual_kind!r}")
+        return LeapAudio(mono, sample_rate, _waveform_graph(mono, duration))
 
     v = np.asarray(visual)
 
-    if visual_kind == 'waveform':
-        return LeapAudio(mono, sample_rate, _waveform_graph(v, duration), x_range)
+    if v.ndim == 1:  # waveform
+        return LeapAudio(mono, sample_rate, _waveform_graph(v, duration))
+
+    if v.ndim == 3:  # already-rendered image
+        if v.shape[-1] not in (1, 3):
+            raise ValueError(f'image visual must have 1 (grayscale) or 3 (rgb) channels, got shape {v.shape}')
+        return LeapAudio(mono, sample_rate, LeapImage(v))
 
-    if visual_kind == 'image':
-        return LeapAudio(mono, sample_rate, LeapImage(v), x_range)
+    if v.ndim == 2:  # spectrogram
+        spec = v
+        if time_axis_index % 2 == 0:  # move time onto the horizontal axis (columns)
+            spec = spec.T
+        img = rescale_min_max(spec.astype(np.float32))[..., None]  # [freq, time, 1]
+        return LeapAudio(mono, sample_rate, LeapImage(img))
 
-    # spectrogram
-    spec = np.squeeze(v)
-    if spec.ndim != 2:
-        raise ValueError(f'spectrogram visual must be 2-D (freq x time), got shape {v.shape}')
-    if time_axis_index % spec.ndim == 0:  # move time onto the horizontal axis (columns)
-        spec = spec.T
-    img = rescale_min_max(spec.astype(np.float32))[..., None]  # [freq, time, 1]
-    return LeapAudio(mono, sample_rate, LeapImage(img), x_range)
+    raise ValueError(f'visual must be 1-D, 2-D, or 3-D, got shape {v.shape}')
 
 
 def default_graph_visualizer(data: npt.NDArray[np.float32]) -> LeapGraph:

{code_loader-1.0.199 → code_loader-1.0.201}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "code-loader"
-version = "1.0.199"
+version = "1.0.201"
 description = ""
 authors = ["dorhar <doron.harnoy@tensorleap.ai>"]
 license = "MIT"