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

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

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

{code_loader-1.0.197.dev0 → code_loader-1.0.197.dev1}/code_loader/contract/datasetclasses.py

@@ -252,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.

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

@@ -514,7 +514,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 +532,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 +545,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:
@@ -906,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 "

{code_loader-1.0.197.dev0 → code_loader-1.0.197.dev1}/code_loader/inner_leap_binder/leapbinder_decorators.py

@@ -1621,7 +1621,11 @@ def tensorleap_metadata(
     return decorating_function
 
 
-def tensorleap_custom_latent_space(name: Optional[str] = None):
+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], preprocess_response: PreprocessResponse):
@@ -1664,7 +1668,8 @@ def tensorleap_custom_latent_space(name: Optional[str] = None):
 
             return result
 
-        leap_binder.set_custom_latent_space(inner_without_validate, ls_name)
+        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.197.dev0 → code_loader-1.0.197.dev1}/code_loader/leaploader.py

@@ -1076,6 +1076,19 @@ class LeapLoader(LeapLoaderBase):
         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:
         self.exec_script()

{code_loader-1.0.197.dev0 → code_loader-1.0.197.dev1}/code_loader/leaploaderbase.py

@@ -222,6 +222,10 @@ class LeapLoaderBase:
     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.197.dev0 → code_loader-1.0.197.dev1}/pyproject.toml

@@ -1,7 +1,7 @@
 [tool.poetry]
 name = "code-loader"
 
-version = "1.0.197.dev0"
+version = "1.0.197.dev1"
 description = ""
 authors = ["dorhar <doron.harnoy@tensorleap.ai>"]
 license = "MIT"