code-loader 1.0.183.dev0__tar.gz → 1.0.184.dev0__tar.gz

{code_loader-1.0.183.dev0 → code_loader-1.0.184.dev0}/PKG-INFO

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

{code_loader-1.0.183.dev0 → code_loader-1.0.184.dev0}/code_loader/leaploader.py

@@ -44,7 +44,7 @@ class LeapLoader(LeapLoaderBase):
         super().__init__(code_path, code_entry_name)
 
         self._preprocess_result_cached = None
-        self._synthetic_responses: Dict[str, PreprocessResponse] = {}
+        self._synthetic_lookup: Dict[str, Tuple[PreprocessResponse, Any]] = {}
 
         try:
             from code_loader.mixpanel_tracker import track_code_loader_loaded
@@ -191,22 +191,11 @@ class LeapLoader(LeapLoaderBase):
     def get_sample(self, state: DataStateEnum, sample_id: Union[int, str], instance_id: int = None) -> DatasetSample:
         self.exec_script()
 
-        if isinstance(sample_id, str) and sample_id.startswith("synthetic_"):
-            last_underscore = sample_id.rfind("_")
-            if last_underscore != -1:
-                prefix = sample_id[:last_underscore]
-                local_idx_str = sample_id[last_underscore + 1:]
-                sim_preprocess = self._synthetic_responses.get(prefix)
-                if sim_preprocess is not None:
-                    try:
-                        local_idx = int(local_idx_str)
-                        original_sample_id = sim_preprocess.sample_ids[local_idx]
-                    except (ValueError, IndexError):
-                        pass
-                    else:
-                        return self._get_sample_from_preprocess(
-                            sim_preprocess, original_sample_id, synthetic_index=sample_id
-                        )
+        if isinstance(sample_id, str) and sample_id in self._synthetic_lookup:
+            sim_preprocess, original_local_id = self._synthetic_lookup[sample_id]
+            return self._get_sample_from_preprocess(
+                sim_preprocess, original_local_id, synthetic_index=sample_id
+            )
 
         preprocess_result = self._preprocess_result()
         if state == DataStateEnum.unlabeled and sample_id not in preprocess_result[state].sample_ids:
@@ -408,8 +397,8 @@ class LeapLoader(LeapLoaderBase):
             result_payloads.append(test_result)
         return result_payloads
 
-    def run_simulation(self, sim_name, params=None, n_samples=1, seed=0, sample_id_prefix=None):
-        # type: (str, Optional[Dict[str, Any]], int, int, Optional[str]) -> Dict[str, Any]
+    def run_simulation(self, sim_name, params=None, n_samples=1, seed=0, sample_ids=None):
+        # type: (str, Optional[Dict[str, Any]], int, int, Optional[List[str]]) -> Dict[str, Any]
         self.exec_script()
         sim = next(
             (s for s in global_leap_binder.setup_container.simulations if s.name == sim_name),
@@ -433,13 +422,43 @@ class LeapLoader(LeapLoaderBase):
             for handler in global_leap_binder.setup_container.inputs:
                 per_encoder[handler.name].append(handler.function(sample_id, sim_preprocess))
         encoded = {name: np.stack(arrays) for name, arrays in per_encoder.items()}
-        if sample_id_prefix is not None:
-            returned_sample_ids = ["{}_{}".format(sample_id_prefix, i) for i in range(len(original_sample_ids))]
-            self._synthetic_responses[sample_id_prefix] = sim_preprocess
+        if sample_ids is not None:
+            if len(sample_ids) != len(original_sample_ids):
+                raise ValueError(
+                    "sample_ids length ({}) does not match simulation output length ({})".format(
+                        len(sample_ids), len(original_sample_ids)
+                    )
+                )
+            for sid in sample_ids:
+                if not isinstance(sid, str):
+                    raise TypeError(
+                        "All sample_ids must be of type str. Got: {}".format(type(sid))
+                    )
+            for synth_id, original_local_id in zip(sample_ids, original_sample_ids):
+                self._synthetic_lookup[synth_id] = (sim_preprocess, original_local_id)
+            self._extend_additional_preprocess(list(sample_ids))
+            returned_sample_ids = list(sample_ids)
         else:
             returned_sample_ids = original_sample_ids
         return {"encoded": encoded, "sample_ids": returned_sample_ids}
 
+    def _extend_additional_preprocess(self, new_sample_ids: List[str]) -> None:
+        if self._preprocess_result_cached is None:
+            self._preprocess_result()
+        additional = self._preprocess_result_cached.get(DataStateEnum.additional)
+        if additional is None:
+            placeholder = PreprocessResponse(
+                sample_ids=list(new_sample_ids),
+                data={},
+                state=DataStateType.additional,
+                sample_id_type=str,
+            )
+            placeholder.tl_generated = True
+            self._preprocess_result_cached[DataStateEnum.additional] = placeholder
+        else:
+            additional.sample_ids = list(additional.sample_ids) + list(new_sample_ids)
+            additional.length = len(additional.sample_ids)
+
     @staticmethod
     def _get_all_dataset_base_handlers() -> List[Union[DatasetBaseHandler, MetadataHandler]]:
         all_dataset_base_handlers: List[Union[DatasetBaseHandler, MetadataHandler]] = []
@@ -653,7 +672,13 @@ class LeapLoader(LeapLoaderBase):
                               state: DataStateEnum, sample_id: Union[int, str]) -> Dict[str, npt.NDArray[np.float32]]:
         result_agg = {}
         preprocess_result = self._preprocess_result()
-        preprocess_state = preprocess_result[state]
+        if (state == DataStateEnum.additional and isinstance(sample_id, str)
+                and sample_id in self._synthetic_lookup):
+            sim_preprocess, original_local_id = self._synthetic_lookup[sample_id]
+            preprocess_state = sim_preprocess
+            sample_id = original_local_id
+        else:
+            preprocess_state = preprocess_result[state]
         for handler in handlers:
             handler_result = handler.function(sample_id, preprocess_state)
             handler_name = handler.name
@@ -667,7 +692,13 @@ class LeapLoader(LeapLoaderBase):
         if instance_id is None:
             return None
         preprocess_result = self._preprocess_result()
-        preprocess_state = preprocess_result[state]
+        if (state == DataStateEnum.additional and isinstance(sample_id, str)
+                and sample_id in self._synthetic_lookup):
+            sim_preprocess, original_local_id = self._synthetic_lookup[sample_id]
+            preprocess_state = sim_preprocess
+            sample_id = original_local_id
+        else:
+            preprocess_state = preprocess_result[state]
         result_agg = {}
         for handler in global_leap_binder.setup_container.instance_masks:
             handler_result = handler.function(sample_id, preprocess_state, instance_id)
@@ -727,7 +758,13 @@ class LeapLoader(LeapLoaderBase):
         result_agg = {}
         is_none = {}
         preprocess_result = self._preprocess_result()
-        preprocess_state = preprocess_result[state]
+        if (state == DataStateEnum.additional and isinstance(sample_id, str)
+                and sample_id in self._synthetic_lookup):
+            sim_preprocess, original_local_id = self._synthetic_lookup[sample_id]
+            preprocess_state = sim_preprocess
+            sample_id = original_local_id
+        else:
+            preprocess_state = preprocess_result[state]
         for handler in global_leap_binder.setup_container.metadata:
             if requested_metadata_names:
                 if not is_metadata_name_starts_with_handler_name(handler):

{code_loader-1.0.183.dev0 → code_loader-1.0.184.dev0}/code_loader/leaploaderbase.py

@@ -154,8 +154,8 @@ class LeapLoaderBase:
         pass
 
     @abstractmethod
-    def run_simulation(self, sim_name, params=None, n_samples=1, seed=0, sample_id_prefix=None):
-        # type: (str, Optional[Dict[str, Any]], int, int, Optional[str]) -> Dict[str, Any]
+    def run_simulation(self, sim_name, params=None, n_samples=1, seed=0, sample_ids=None):
+        # type: (str, Optional[Dict[str, Any]], int, int, Optional[List[str]]) -> Dict[str, Any]
         pass
 
     def is_custom_latent_space(self) -> bool:

{code_loader-1.0.183.dev0 → code_loader-1.0.184.dev0}/pyproject.toml

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