code-loader 1.0.190.dev0__py3-none-any.whl → 1.0.191.dev1__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -1,6 +1,6 @@
1
1
  import warnings
2
2
  from dataclasses import dataclass, field
3
- from typing import Any, Callable, List, Optional, Dict, Union, Type, Literal
3
+ from typing import Any, Callable, List, Optional, Dict, Union, Type, Literal, cast
4
4
  import re
5
5
  import numpy as np
6
6
  import numpy.typing as npt
@@ -15,6 +15,8 @@ custom_latent_space_attribute = "custom_latent_space"
15
15
 
16
16
  _simulation_context: Dict[str, bool] = {"active": False}
17
17
 
18
+ SampleId = Union[int, str]
19
+
18
20
 
19
21
  @dataclass
20
22
  class PreprocessResponse:
@@ -39,27 +41,44 @@ class PreprocessResponse:
39
41
  """
40
42
  length: Optional[int] = None # Deprecated. Please use sample_ids instead
41
43
  data: Any = None
42
- sample_ids: Optional[Union[List[str], List[int]]] = None
44
+ sample_ids: Optional[Union[List[SampleId], List[List[SampleId]]]] = None
43
45
  state: Optional[DataStateType] = None
44
46
  sample_id_type: Optional[Union[Type[str], Type[int]]] = None
45
47
  sample_ids_to_instance_mappings: Optional[Dict[str, List[str]]] = None # in use only for element instance
46
48
  instance_to_sample_ids_mappings: Optional[Dict[str, str]] = None # in use only for element instance
47
49
  tl_generated: bool = False
50
+ _grouped: bool = field(default=False, init=False, repr=False, compare=False)
48
51
 
49
52
  def __post_init__(self) -> None:
50
53
  assert self.sample_ids_to_instance_mappings is None, f"Keep sample_ids_to_instance_mappings None when initializing PreprocessResponse"
51
54
  assert self.instance_to_sample_ids_mappings is None, f"Keep instance_to_sample_ids_mappings None when initializing PreprocessResponse"
52
55
 
53
56
  if self.length is not None and self.sample_ids is None:
54
- self.sample_ids = [i for i in range(self.length)]
57
+ self.sample_ids = cast(List[SampleId], [i for i in range(self.length)])
55
58
  self.sample_id_type = int
59
+ self._grouped = False
56
60
  elif self.length is None and self.sample_ids is not None:
57
- self.length = len(self.sample_ids)
58
- if self.sample_id_type is None:
59
- self.sample_id_type = str
60
- if self.sample_id_type == str:
61
- for sample_id in self.sample_ids:
62
- assert isinstance(sample_id, str), f"Sample id should be of type str. Got: {type(sample_id)}"
61
+ self._grouped = len(self.sample_ids) > 0 and isinstance(self.sample_ids[0], list)
62
+ if self._grouped:
63
+ groups = cast(List[List[SampleId]], self.sample_ids)
64
+ assert len(groups) >= 1, "Grouped PreprocessResponse must have at least one group."
65
+ for group in groups:
66
+ assert isinstance(group, list) and len(group) > 0, \
67
+ "Each group in a grouped PreprocessResponse must be a non-empty list."
68
+ self.sample_id_type = type(groups[0][0])
69
+ for group in groups:
70
+ for sample_id in group:
71
+ assert isinstance(sample_id, self.sample_id_type), \
72
+ f"Sample id should be of type {self.sample_id_type.__name__}. Got: {type(sample_id)}"
73
+ self.length = sum(len(group) for group in groups)
74
+ else:
75
+ flat_ids = cast(List[SampleId], self.sample_ids)
76
+ self.length = len(flat_ids)
77
+ if self.sample_id_type is None:
78
+ self.sample_id_type = str
79
+ if self.sample_id_type == str:
80
+ for sample_id in flat_ids:
81
+ assert isinstance(sample_id, str), f"Sample id should be of type str. Got: {type(sample_id)}"
63
82
  else:
64
83
  raise Exception("length is deprecated, please use sample_ids instead.")
65
84
 
@@ -79,9 +98,36 @@ class PreprocessResponse:
79
98
  return id(self)
80
99
 
81
100
  def __len__(self) -> int:
101
+ assert self.length is not None
102
+ return self.length
103
+
104
+ @property
105
+ def is_grouped(self) -> bool:
106
+ return self._grouped
107
+
108
+ @property
109
+ def num_groups(self) -> int:
110
+ # Grouped: the number of groups. Flat: the (flat) sample count, since sample_ids is the
111
+ # flat id list and there is one sample per "group".
82
112
  assert self.sample_ids is not None
83
113
  return len(self.sample_ids)
84
114
 
115
+ @property
116
+ def groups(self) -> Optional[List[List[SampleId]]]:
117
+ # The nested list of groups when grouped, otherwise None. Callers should
118
+ # branch on is_grouped before relying on this.
119
+ return cast(List[List[SampleId]], self.sample_ids) if self._grouped else None
120
+
121
+ @property
122
+ def flat_sample_ids(self) -> List[SampleId]:
123
+ # All sample ids in order: concatenation of the groups when grouped,
124
+ # otherwise the flat sample_ids list as-is.
125
+ assert self.sample_ids is not None
126
+ if self._grouped:
127
+ groups = cast(List[List[SampleId]], self.sample_ids)
128
+ return [sample_id for group in groups for sample_id in group]
129
+ return cast(List[SampleId], self.sample_ids)
130
+
85
131
  @dataclass
86
132
  class ElementInstance:
87
133
  name: str
@@ -606,6 +606,15 @@ class LeapBinder:
606
606
  if state_enum in preprocess_result_dict:
607
607
  preprocess_result_dict_in_correct_order[state_enum] = preprocess_result_dict[state_enum]
608
608
 
609
+ # All-or-none: every state must agree on grouped (nested sample_ids) vs flat.
610
+ if len({r.is_grouped for r in preprocess_result_dict_in_correct_order.values()}) > 1:
611
+ shapes = ', '.join(
612
+ f"{state.name}={'grouped' if r.is_grouped else 'flat'}"
613
+ for state, r in preprocess_result_dict_in_correct_order.items())
614
+ raise Exception(
615
+ "All preprocess states must be either all grouped (nested sample_ids) or all flat, "
616
+ f"but got a mix: {shapes}.")
617
+
609
618
  return preprocess_result_dict_in_correct_order
610
619
 
611
620
  def get_preprocess_unlabeled_result(self) -> Optional[PreprocessResponse]:
@@ -626,7 +635,13 @@ class LeapBinder:
626
635
  preprocess_response: PreprocessResponse, test_result: List[DatasetTestResultPayload],
627
636
  dataset_base_handler: Union[DatasetBaseHandler, MetadataHandler], state: DataStateEnum) -> List[DatasetTestResultPayload]:
628
637
  assert preprocess_response.sample_ids is not None
629
- raw_result = dataset_base_handler.function(preprocess_response.sample_ids[0], preprocess_response)
638
+ if preprocess_response.is_grouped:
639
+ # Grouped: probe with the first group, then reduce to a single sample's result so the
640
+ # recorded shape/type is per-sample (matching the flat case), not the (B, *) batch.
641
+ group = preprocess_response.sample_ids[0]
642
+ raw_result = dataset_base_handler.function(group, preprocess_response)[0]
643
+ else:
644
+ raw_result = dataset_base_handler.function(preprocess_response.sample_ids[0], preprocess_response)
630
645
  handler_type = 'metadata' if isinstance(dataset_base_handler, MetadataHandler) else None
631
646
  if isinstance(dataset_base_handler, MetadataHandler):
632
647
  if isinstance(raw_result, dict):
@@ -32,6 +32,7 @@ from code_loader.mixpanel_tracker import clear_integration_events, AnalyticsEven
32
32
 
33
33
  _called_from_inside_tl_decorator = 0
34
34
  _called_from_inside_tl_integration_test_decorator = False
35
+ _mapping_dataset_is_grouped = False
35
36
  _call_from_tl_platform = os.environ.get('IS_TENSORLEAP_PLATFORM') == 'true'
36
37
 
37
38
  # ---- warnings store (module-level) ----
@@ -197,14 +198,58 @@ def validate_output_structure(result, func_name: str, expected_type_name="np.nda
197
198
  def batch_warning(result, func_name):
198
199
  if len(result.shape) > 0 and result.shape[0] == 1:
199
200
  warnings.warn(
200
- f"{func_name} warning: Tensorleap will add a batch dimension at axis 0 of "
201
- f"{func_name}'s output, but axis 0 is already size 1 — this can produce a "
202
- f"doubly-batched shape (e.g. (1, 1, ...)). If axis 0 is a real dimension "
203
- f"(e.g. a single channel) and not an accidental batch, you can ignore this; "
204
- f"only fix it if you pre-batched the output."
201
+ f"{func_name} warning: Tensorleap will add a batch dimension at axis 0 to the output of {func_name}, "
202
+ f"although the detected size of axis 0 is already 1. "
203
+ f"This may lead to an extra batch dimension (e.g., shape (1, 1, ...)). "
204
+ f"Please ensure that the output of '{func_name}' is not already batched "
205
+ f"to avoid computation errors."
205
206
  )
206
207
 
207
208
 
209
+ def _validate_id_or_group(sample_id, preprocess_response, func_name):
210
+ """Accept either a single sample id (flat dataset) or a group — a list of ids (grouped
211
+ dataset). The argument shape must match the dataset: a group requires a grouped
212
+ PreprocessResponse and a single id requires a flat one. Every id must match the declared
213
+ sample_id_type."""
214
+ is_group = isinstance(sample_id, list)
215
+ if is_group and not preprocess_response.is_grouped:
216
+ raise AssertionError(
217
+ f'{func_name}() validation failed: got a group (list of sample ids) but the '
218
+ f'PreprocessResponse is not grouped. Pass a single sample id for a flat dataset.')
219
+ if not is_group and preprocess_response.is_grouped:
220
+ raise AssertionError(
221
+ f'{func_name}() validation failed: got a single sample id but the PreprocessResponse '
222
+ f'is grouped. Pass a group (list of sample ids) for a grouped dataset.')
223
+ ids = sample_id if is_group else [sample_id]
224
+ for sid in ids:
225
+ assert type(sid) == preprocess_response.sample_id_type, \
226
+ (f'{func_name}() validation failed: '
227
+ f'Argument sample_id should be as the same type as defined in the preprocess response '
228
+ f'{preprocess_response.sample_id_type}. Got {type(sid)}.')
229
+
230
+
231
+ def _validate_grouped_result(result, group_size, func_name, validate_single):
232
+ """A grouped (group) encoder result must be either a single ndarray with leading dim
233
+ == group_size, or a list of `group_size` per-sample results (each validated by
234
+ validate_single). Metadata validates its grouped result separately (list only)."""
235
+ if isinstance(result, np.ndarray):
236
+ assert len(result.shape) > 0 and result.shape[0] == group_size, \
237
+ (f'{func_name}() validation failed: expected a grouped output for a group of '
238
+ f'{group_size}: an array with leading dim {group_size}, got shape {result.shape}.')
239
+ validate_single(result)
240
+ elif isinstance(result, list):
241
+ assert len(result) == group_size, \
242
+ (f'{func_name}() validation failed: expected a grouped output for a group of '
243
+ f'{group_size}: a list of {group_size} arrays, got {len(result)}.')
244
+ for r in result:
245
+ validate_single(r)
246
+ else:
247
+ raise AssertionError(
248
+ f'{func_name}() validation failed: expected a grouped output for a group of '
249
+ f'{group_size}: either an array with leading dim {group_size} or a list of '
250
+ f'{group_size} arrays, got {type(result)}.')
251
+
252
+
208
253
  def _add_mapping_connection(user_unique_name, connection_destinations, arg_names, name, node_mapping_type):
209
254
  connection_destinations = [connection_destination for connection_destination in connection_destinations
210
255
  if not isinstance(connection_destination, SamplePreprocessResponse)]
@@ -282,16 +327,12 @@ def tensorleap_integration_test():
282
327
 
283
328
  def _validate_input_args(*args, **kwargs):
284
329
  sample_id, preprocess_response = args
285
- assert type(sample_id) == preprocess_response.sample_id_type, (
286
- f"tensorleap_integration_test validation failed: "
287
- f"sample_id type ({type(sample_id).__name__}) does not match the expected "
288
- f"type ({preprocess_response.sample_id_type}) from the PreprocessResponse."
289
- )
330
+ _validate_id_or_group(sample_id, preprocess_response, 'tensorleap_integration_test')
290
331
 
291
332
  def inner(*args, **kwargs):
292
333
  if not _call_from_tl_platform:
293
334
  set_current('tensorleap_integration_test')
294
- validate_args_structure(*args, types_order=[Union[int, str], PreprocessResponse],
335
+ validate_args_structure(*args, types_order=[Union[int, str, list], PreprocessResponse],
295
336
  func_name='integration_test', expected_names=["idx", "preprocess"], **kwargs)
296
337
  _validate_input_args(*args, **kwargs)
297
338
 
@@ -308,6 +349,8 @@ def tensorleap_integration_test():
308
349
  "v") # put here because otherwise it will become v only if it finishes all the script
309
350
  ret = integration_test_function(*args, **kwargs)
310
351
 
352
+ global _mapping_dataset_is_grouped
353
+ _mapping_dataset_is_grouped = args[1].is_grouped
311
354
  try:
312
355
  os.environ[mapping_runtime_mode_env_var_mame] = 'True'
313
356
  integration_test_function(None, PreprocessResponse(state=DataStateType.training, length=0))
@@ -318,15 +361,16 @@ def tensorleap_integration_test():
318
361
  line_number = first_tb.lineno
319
362
  if isinstance(e, TypeError) and 'is not subscriptable' in str(e):
320
363
  update_env_params_func("code_mapping", "x")
321
- raise (f'Invalid integration code. File {file_name}, line {line_number}: '
364
+ raise Exception(f'Invalid integration code. File {file_name}, line {line_number}: '
322
365
  f"indexing is supported only on the model's predictions inside the integration test. Please remove this indexing operation usage from the integration test code.")
323
366
  else:
324
367
  update_env_params_func("code_mapping", "x")
325
368
 
326
- raise (f'Invalid integration code. File {file_name}, line {line_number}: '
369
+ raise Exception(f'Invalid integration code. File {file_name}, line {line_number}: '
327
370
  f'Integration test is only allowed to call Tensorleap decorators. '
328
371
  f'Ensure any arithmetics, external library use, Python logic is placed within Tensorleap decoders')
329
372
  finally:
373
+ _mapping_dataset_is_grouped = False
330
374
  if mapping_runtime_mode_env_var_mame in os.environ:
331
375
  del os.environ[mapping_runtime_mode_env_var_mame]
332
376
  finally:
@@ -349,14 +393,6 @@ def _safe_get_item(key):
349
393
 
350
394
 
351
395
  def tensorleap_load_model(prediction_types: Optional[List[PredictionTypeHandler]] = _UNSET):
352
- """Register the model-loading function.
353
-
354
- ``prediction_types``: declare one ``PredictionTypeHandler`` per model output, in
355
- output order. Multi-output models must declare an entry for EVERY output — the
356
- count is validated against the model's outputs — so for an output you don't consume
357
- as a prediction, pass a throwaway handler to satisfy the count. Omit the argument
358
- only when the model has a single prediction output.
359
- """
360
396
  prediction_types_was_provided = prediction_types is not _UNSET
361
397
 
362
398
  if not prediction_types_was_provided:
@@ -575,9 +611,6 @@ def tensorleap_load_model(prediction_types: Optional[List[PredictionTypeHandler]
575
611
  def get_inputs(self):
576
612
  return self.model.get_inputs()
577
613
 
578
- def get_outputs(self):
579
- return self.model.get_outputs()
580
-
581
614
  model_placeholder = ModelPlaceholder(prediction_types)
582
615
  if not _call_from_tl_platform:
583
616
  update_env_params_func("tensorleap_load_model", "v")
@@ -642,13 +675,6 @@ def tensorleap_load_model(prediction_types: Optional[List[PredictionTypeHandler]
642
675
 
643
676
  return FollowInputIndex()
644
677
 
645
- def get_outputs(self):
646
- # Mapping-mode counterpart of the ONNX-run get_outputs(): in mapping
647
- # mode outputs are reached by index, and get_outputs()[i] resolves to
648
- # Prediction{i} — exactly like run()/__call__, which also return a
649
- # ModelOutputPlaceholder. Reusing it keeps the mapping unchanged.
650
- return ModelOutputPlaceholder()
651
-
652
678
  return ModelPlaceholder()
653
679
 
654
680
  def final_inner(*args, **kwargs):
@@ -1526,13 +1552,10 @@ def tensorleap_metadata(
1526
1552
  raise Exception(f'Metadata with name {name} already exists. '
1527
1553
  f'Please choose another')
1528
1554
 
1529
- def _validate_input_args(sample_id: Union[int, str], preprocess_response: PreprocessResponse):
1530
- assert type(sample_id) == preprocess_response.sample_id_type, \
1531
- (f'{user_function.__name__}() validation failed: '
1532
- f'Argument sample_id should be as the same type as defined in the preprocess response '
1533
- f'{preprocess_response.sample_id_type}. Got {type(sample_id)}.')
1555
+ def _validate_input_args(sample_id: Union[int, str, list], preprocess_response: PreprocessResponse):
1556
+ _validate_id_or_group(sample_id, preprocess_response, user_function.__name__)
1534
1557
 
1535
- def _validate_result(result):
1558
+ def _validate_single(result):
1536
1559
  supported_result_types = (type(None), int, str, bool, float, dict, np.floating,
1537
1560
  np.bool_, np.unsignedinteger, np.signedinteger, np.integer)
1538
1561
  if isinstance(result, tuple):
@@ -1550,6 +1573,20 @@ def tensorleap_metadata(
1550
1573
  (f'{user_function.__name__}() validation failed: '
1551
1574
  f'Values in the return dict should be of type {str(supported_result_types)}. Got {type(value)}.')
1552
1575
 
1576
+ def _validate_result(result, grouped=False, group_size=None):
1577
+ if not grouped:
1578
+ _validate_single(result)
1579
+ return
1580
+ # Grouped metadata: a list of `group_size` per-sample scalars/dicts.
1581
+ assert isinstance(result, list), \
1582
+ (f'{user_function.__name__}() validation failed: expected a grouped output for a group '
1583
+ f'of {group_size}: a list of {group_size} metadata values, got {type(result)}.')
1584
+ assert len(result) == group_size, \
1585
+ (f'{user_function.__name__}() validation failed: expected a grouped output for a group '
1586
+ f'of {group_size}: a list of {group_size} metadata values, got {len(result)}.')
1587
+ for value in result:
1588
+ _validate_single(value)
1589
+
1553
1590
  def inner_without_validate(sample_id, preprocess_response):
1554
1591
 
1555
1592
  global _called_from_inside_tl_decorator
@@ -1569,14 +1606,16 @@ def tensorleap_metadata(
1569
1606
  set_current('tensorleap_metadata')
1570
1607
  if os.environ.get(mapping_runtime_mode_env_var_mame):
1571
1608
  return None
1572
- validate_args_structure(*args, types_order=[Union[int, str], PreprocessResponse],
1609
+ validate_args_structure(*args, types_order=[Union[int, str, list], PreprocessResponse],
1573
1610
  func_name=user_function.__name__, expected_names=["idx", "preprocess"], **kwargs)
1574
1611
  sample_id, preprocess_response = args if len(args) != 0 else kwargs.values()
1575
1612
  _validate_input_args(sample_id, preprocess_response)
1576
1613
 
1614
+ grouped = isinstance(sample_id, list)
1615
+ group_size = len(sample_id) if grouped else None
1577
1616
  result = inner_without_validate(sample_id, preprocess_response)
1578
1617
 
1579
- _validate_result(result)
1618
+ _validate_result(result, grouped, group_size)
1580
1619
  if not _call_from_tl_platform:
1581
1620
  update_env_params_func("tensorleap_metadata", "v")
1582
1621
  return result
@@ -1588,35 +1627,50 @@ def tensorleap_metadata(
1588
1627
 
1589
1628
  def tensorleap_custom_latent_space():
1590
1629
  def decorating_function(user_function: SectionCallableInterface):
1591
- def _validate_input_args(sample_id: Union[int, str], preprocess_response: PreprocessResponse):
1592
- assert isinstance(sample_id, (int, str)), \
1593
- (f'tensorleap_custom_latent_space validation failed: '
1594
- f'Argument sample_id should be either int or str. Got {type(sample_id)}.')
1595
- assert isinstance(preprocess_response, PreprocessResponse), \
1596
- (f'tensorleap_custom_latent_space validation failed: '
1597
- f'Argument preprocess_response should be a PreprocessResponse. Got {type(preprocess_response)}.')
1598
- assert type(sample_id) == preprocess_response.sample_id_type, \
1599
- (f'tensorleap_custom_latent_space validation failed: '
1600
- f'Argument sample_id should be as the same type as defined in the preprocess response '
1601
- f'{preprocess_response.sample_id_type}. Got {type(sample_id)}.')
1630
+ def _validate_input_args(sample_id: Union[int, str, list], preprocess_response: PreprocessResponse):
1631
+ _validate_id_or_group(sample_id, preprocess_response, 'tensorleap_custom_latent_space')
1602
1632
 
1603
- def _validate_result(result):
1604
- assert isinstance(result, np.ndarray), \
1633
+ def _validate_single(single_result):
1634
+ assert isinstance(single_result, np.ndarray), \
1605
1635
  (f'tensorleap_custom_latent_space validation failed: '
1606
- f'The return type should be a numpy array. Got {type(result)}.')
1607
- if result.ndim > 1:
1608
- flat_dim = int(np.prod(result.shape))
1636
+ f'The return type should be a numpy array. Got {type(single_result)}.')
1637
+ if single_result.ndim > 1:
1638
+ flat_dim = int(np.prod(single_result.shape))
1609
1639
  store_general_warning(
1610
- key=("tensorleap_custom_latent_space_flatten", tuple(result.shape)),
1640
+ key=("tensorleap_custom_latent_space_flatten", tuple(single_result.shape)),
1611
1641
  message=(
1612
- f"tensorleap_custom_latent_space returned per-sample shape {tuple(result.shape)} "
1613
- f"(ndim={result.ndim}). Tensorleap assumes per-sample shape (d, ...) and will "
1642
+ f"tensorleap_custom_latent_space returned per-sample shape {tuple(single_result.shape)} "
1643
+ f"(ndim={single_result.ndim}). Tensorleap assumes per-sample shape (d, ...) and will "
1614
1644
  f"flatten to ({flat_dim},) before downstream visualization and clustering. "
1615
1645
  f"If you want a different aggregation (e.g. global average pooling), do it "
1616
1646
  f"inside your function."
1617
1647
  ),
1618
1648
  )
1619
1649
 
1650
+ def _validate_result(result, grouped=False, group_size=None):
1651
+ if not grouped:
1652
+ _validate_single(result)
1653
+ return
1654
+ # Grouped: (B, *) array or list of B per-sample arrays. Validate each per-sample slice so
1655
+ # the per-sample flatten check/warning uses the sample shape, not the (B, *) batch shape.
1656
+ if isinstance(result, np.ndarray):
1657
+ assert result.ndim >= 1 and result.shape[0] == group_size, \
1658
+ (f'tensorleap_custom_latent_space validation failed: expected a grouped output for a '
1659
+ f'group of {group_size}: an array with leading dim {group_size}, got shape {result.shape}.')
1660
+ for single in result:
1661
+ _validate_single(single)
1662
+ elif isinstance(result, list):
1663
+ assert len(result) == group_size, \
1664
+ (f'tensorleap_custom_latent_space validation failed: expected a grouped output for a '
1665
+ f'group of {group_size}: a list of {group_size} arrays, got {len(result)}.')
1666
+ for single in result:
1667
+ _validate_single(single)
1668
+ else:
1669
+ raise AssertionError(
1670
+ f'tensorleap_custom_latent_space validation failed: expected a grouped output for a '
1671
+ f'group of {group_size}: either an array with leading dim {group_size} or a list of '
1672
+ f'{group_size} arrays, got {type(result)}.')
1673
+
1620
1674
  def inner_without_validate(sample_id, preprocess_response):
1621
1675
  global _called_from_inside_tl_decorator
1622
1676
  _called_from_inside_tl_decorator += 1
@@ -1636,9 +1690,11 @@ def tensorleap_custom_latent_space():
1636
1690
 
1637
1691
  _validate_input_args(sample_id, preprocess_response)
1638
1692
 
1693
+ grouped = isinstance(sample_id, list)
1694
+ group_size = len(sample_id) if grouped else None
1639
1695
  result = inner_without_validate(sample_id, preprocess_response)
1640
1696
 
1641
- _validate_result(result)
1697
+ _validate_result(result, grouped, group_size)
1642
1698
  return result
1643
1699
 
1644
1700
  return inner
@@ -1676,6 +1732,15 @@ def tensorleap_preprocess():
1676
1732
  assert len(set(result)) == len(result), \
1677
1733
  (f'{user_function.__name__}() validation failed: '
1678
1734
  f'The return list should not contain duplicate PreprocessResponse objects.')
1735
+ # All-or-none: a dataset is either all grouped (nested sample_ids) or all flat.
1736
+ if len({response.is_grouped for response in result}) != 1:
1737
+ shapes = ', '.join(
1738
+ f"#{i}={'grouped' if response.is_grouped else 'flat'}"
1739
+ for i, response in enumerate(result))
1740
+ raise AssertionError(
1741
+ f'{user_function.__name__}() validation failed: '
1742
+ f'All PreprocessResponses must be either all grouped (nested sample_ids) '
1743
+ f'or all flat, but got a mix: {shapes}.')
1679
1744
 
1680
1745
  def inner(*args, **kwargs):
1681
1746
  if not _call_from_tl_platform:
@@ -2013,14 +2078,10 @@ def tensorleap_input_encoder(name: str, channel_dim=_UNSET, model_input_index=No
2013
2078
  f"channel axis in the batched tensor (batch = axis 0), so 0 is invalid. Use 1 "
2014
2079
  f"for channel-first (NCHW) and -1 for channel-last (NHWC).")
2015
2080
 
2016
- def _validate_input_args(sample_id: Union[int, str], preprocess_response: PreprocessResponse):
2017
- assert type(sample_id) == preprocess_response.sample_id_type, \
2018
- (f'{user_function.__name__}() validation failed: '
2019
- f'Argument sample_id should be as the same type as defined in the preprocess response '
2020
- f'{preprocess_response.sample_id_type}. Got {type(sample_id)}.')
2081
+ def _validate_input_args(sample_id: Union[int, str, list], preprocess_response: PreprocessResponse):
2082
+ _validate_id_or_group(sample_id, preprocess_response, user_function.__name__)
2021
2083
 
2022
- def _validate_result(result):
2023
- validate_output_structure(result, func_name=user_function.__name__, expected_type_name="np.ndarray")
2084
+ def _validate_single(result):
2024
2085
  assert isinstance(result, np.ndarray), \
2025
2086
  (f'{user_function.__name__}() validation failed: '
2026
2087
  f'Unsupported return type. Should be a numpy array. Got {type(result)}.')
@@ -2030,6 +2091,13 @@ def tensorleap_input_encoder(name: str, channel_dim=_UNSET, model_input_index=No
2030
2091
  assert channel_dim - 1 <= len(result.shape), (f'{user_function.__name__}() validation failed: '
2031
2092
  f'The channel_dim ({channel_dim}) should be <= to the rank of the resulting input rank ({len(result.shape)}).')
2032
2093
 
2094
+ def _validate_result(result, grouped=False, group_size=None):
2095
+ if not grouped:
2096
+ validate_output_structure(result, func_name=user_function.__name__, expected_type_name="np.ndarray")
2097
+ _validate_single(result)
2098
+ return
2099
+ _validate_grouped_result(result, group_size, user_function.__name__, _validate_single)
2100
+
2033
2101
  def inner_without_validate(sample_id, preprocess_response):
2034
2102
  global _called_from_inside_tl_decorator
2035
2103
  _called_from_inside_tl_decorator += 1
@@ -2046,16 +2114,18 @@ def tensorleap_input_encoder(name: str, channel_dim=_UNSET, model_input_index=No
2046
2114
  def inner(*args, **kwargs):
2047
2115
  if not _call_from_tl_platform:
2048
2116
  set_current("tensorleap_input_encoder")
2049
- validate_args_structure(*args, types_order=[Union[int, str], PreprocessResponse],
2117
+ validate_args_structure(*args, types_order=[Union[int, str, list], PreprocessResponse],
2050
2118
  func_name=user_function.__name__, expected_names=["idx", "preprocess"], **kwargs)
2051
2119
  sample_id, preprocess_response = args if len(args) != 0 else kwargs.values()
2052
2120
  _validate_input_args(sample_id, preprocess_response)
2053
2121
 
2122
+ grouped = isinstance(sample_id, list)
2123
+ group_size = len(sample_id) if grouped else None
2054
2124
  result = inner_without_validate(sample_id, preprocess_response)
2055
2125
 
2056
- _validate_result(result)
2126
+ _validate_result(result, grouped, group_size)
2057
2127
 
2058
- if _called_from_inside_tl_decorator == 0 and _called_from_inside_tl_integration_test_decorator:
2128
+ if not grouped and _called_from_inside_tl_decorator == 0 and _called_from_inside_tl_integration_test_decorator:
2059
2129
  batch_warning(result, user_function.__name__)
2060
2130
  result = np.expand_dims(result, axis=0)
2061
2131
  # Emit integration test event once per test
@@ -2078,8 +2148,15 @@ def tensorleap_input_encoder(name: str, channel_dim=_UNSET, model_input_index=No
2078
2148
  inner.node_mapping = NodeMapping(name, node_mapping_type)
2079
2149
 
2080
2150
  def mapping_inner(*args, **kwargs):
2081
- class TempMapping:
2082
- pass
2151
+ if _mapping_dataset_is_grouped:
2152
+ class TempMapping:
2153
+ def __getitem__(self, key):
2154
+ # Indexing a grouped input group (input_list[i]) selects a sample from the
2155
+ # same input source; the model wiring is identical to the un-grouped input.
2156
+ return self
2157
+ else:
2158
+ class TempMapping:
2159
+ pass
2083
2160
 
2084
2161
  ret = TempMapping()
2085
2162
  ret.node_mapping = mapping_inner.node_mapping
@@ -2109,15 +2186,10 @@ def tensorleap_gt_encoder(name: str):
2109
2186
  raise Exception(f'GT with name {name} already exists. '
2110
2187
  f'Please choose another')
2111
2188
 
2112
- def _validate_input_args(sample_id: Union[int, str], preprocess_response: PreprocessResponse):
2113
- assert type(sample_id) == preprocess_response.sample_id_type, \
2114
- (f'{user_function.__name__}() validation failed: '
2115
- f'Argument sample_id should be as the same type as defined in the preprocess response '
2116
- f'{preprocess_response.sample_id_type}. Got {type(sample_id)}.')
2189
+ def _validate_input_args(sample_id: Union[int, str, list], preprocess_response: PreprocessResponse):
2190
+ _validate_id_or_group(sample_id, preprocess_response, user_function.__name__)
2117
2191
 
2118
- def _validate_result(result):
2119
- validate_output_structure(result, func_name=user_function.__name__, expected_type_name="np.ndarray",
2120
- gt_flag=True)
2192
+ def _validate_single(result):
2121
2193
  assert isinstance(result, np.ndarray), \
2122
2194
  (f'{user_function.__name__}() validation failed: '
2123
2195
  f'Unsupported return type. Should be a numpy array. Got {type(result)}.')
@@ -2125,6 +2197,14 @@ def tensorleap_gt_encoder(name: str):
2125
2197
  (f'{user_function.__name__}() validation failed: '
2126
2198
  f'The return type should be a numpy array of type float32. Got {result.dtype}.')
2127
2199
 
2200
+ def _validate_result(result, grouped=False, group_size=None):
2201
+ if not grouped:
2202
+ validate_output_structure(result, func_name=user_function.__name__, expected_type_name="np.ndarray",
2203
+ gt_flag=True)
2204
+ _validate_single(result)
2205
+ return
2206
+ _validate_grouped_result(result, group_size, user_function.__name__, _validate_single)
2207
+
2128
2208
  def inner_without_validate(sample_id, preprocess_response):
2129
2209
  global _called_from_inside_tl_decorator
2130
2210
  _called_from_inside_tl_decorator += 1
@@ -2141,16 +2221,18 @@ def tensorleap_gt_encoder(name: str):
2141
2221
  def inner(*args, **kwargs):
2142
2222
  if not _call_from_tl_platform:
2143
2223
  set_current("tensorleap_gt_encoder")
2144
- validate_args_structure(*args, types_order=[Union[int, str], PreprocessResponse],
2224
+ validate_args_structure(*args, types_order=[Union[int, str, list], PreprocessResponse],
2145
2225
  func_name=user_function.__name__, expected_names=["idx", "preprocess"], **kwargs)
2146
2226
  sample_id, preprocess_response = args
2147
2227
  _validate_input_args(sample_id, preprocess_response)
2148
2228
 
2229
+ grouped = isinstance(sample_id, list)
2230
+ group_size = len(sample_id) if grouped else None
2149
2231
  result = inner_without_validate(sample_id, preprocess_response)
2150
2232
 
2151
- _validate_result(result)
2233
+ _validate_result(result, grouped, group_size)
2152
2234
 
2153
- if _called_from_inside_tl_decorator == 0 and _called_from_inside_tl_integration_test_decorator:
2235
+ if not grouped and _called_from_inside_tl_decorator == 0 and _called_from_inside_tl_integration_test_decorator:
2154
2236
  batch_warning(result, user_function.__name__)
2155
2237
  result = np.expand_dims(result, axis=0)
2156
2238
  # Emit integration test event once per test
@@ -2167,8 +2249,15 @@ def tensorleap_gt_encoder(name: str):
2167
2249
  inner.node_mapping = NodeMapping(name, NodeMappingType.GroundTruth)
2168
2250
 
2169
2251
  def mapping_inner(*args, **kwargs):
2170
- class TempMapping:
2171
- pass
2252
+ if _mapping_dataset_is_grouped:
2253
+ class TempMapping:
2254
+ def __getitem__(self, key):
2255
+ # Indexing a grouped GT group (gt_list[i]) selects a sample from the same
2256
+ # ground-truth source; wiring is identical to the un-grouped GT node.
2257
+ return self
2258
+ else:
2259
+ class TempMapping:
2260
+ pass
2172
2261
 
2173
2262
  ret = TempMapping()
2174
2263
  ret.node_mapping = mapping_inner.node_mapping
code_loader/leaploader.py CHANGED
@@ -246,16 +246,20 @@ class LeapLoader(LeapLoaderBase):
246
246
  )
247
247
 
248
248
  preprocess_result = self._preprocess_result()
249
- if state == DataStateEnum.unlabeled and sample_id not in preprocess_result[state].sample_ids:
249
+ if state == DataStateEnum.unlabeled and sample_id not in preprocess_result[state].flat_sample_ids:
250
250
  self._preprocess_result(update_unlabeled_preprocess=True)
251
251
 
252
252
  metadata, metadata_is_none = self.get_metadata(state, sample_id)
253
253
 
254
254
  custom_latent_space = None
255
255
  if global_leap_binder.setup_container.custom_latent_space is not None:
256
- custom_latent_space = global_leap_binder.setup_container.custom_latent_space.function(sample_id,
257
- preprocess_result[
258
- state])
256
+ latent_fn = global_leap_binder.setup_container.custom_latent_space.function
257
+ preprocess_state = preprocess_result[state]
258
+ if preprocess_state.is_grouped:
259
+ group_ids, pos = self._locate_group(preprocess_state, sample_id)
260
+ custom_latent_space = self._to_grouped_array(latent_fn(group_ids, preprocess_state))[pos]
261
+ else:
262
+ custom_latent_space = latent_fn(sample_id, preprocess_state)
259
263
  instance_mask = self._get_instances_masks(state, sample_id, instance_id)
260
264
  sample = DatasetSample(inputs=self._get_inputs(state, sample_id),
261
265
  gt=None if state == DataStateEnum.unlabeled else self._get_gt(state, sample_id),
@@ -320,6 +324,8 @@ class LeapLoader(LeapLoaderBase):
320
324
  try:
321
325
  self.exec_script()
322
326
 
327
+ integration_test_test_payload = self._check_integration_test_exists()
328
+ test_payloads.append(integration_test_test_payload)
323
329
  preprocess_test_payload = self._check_preprocess()
324
330
  test_payloads.append(preprocess_test_payload)
325
331
  handlers_test_payloads = self._check_handlers()
@@ -356,6 +362,23 @@ class LeapLoader(LeapLoaderBase):
356
362
  engine_file_contract=EngineFileContract(global_leap_binder.mapping_connections,
357
363
  global_leap_binder.leap_analysis_configuration))
358
364
 
365
+ def _check_integration_test_exists(self) -> DatasetTestResultPayload:
366
+ test_result = DatasetTestResultPayload('integration_test')
367
+ # Only enforced on the Tensorleap platform (i.e. during a push / inspection,
368
+ # where IS_TENSORLEAP_PLATFORM=true). Running code-loader locally or in unit
369
+ # tests must not require the decorator, so pre-decorator integrations and
370
+ # local check runs keep working unchanged.
371
+ is_on_platform = os.environ.get('IS_TENSORLEAP_PLATFORM') == 'true'
372
+ if is_on_platform and global_leap_binder.integration_test_func is None:
373
+ test_result.is_passed = False
374
+ test_result.display[TestingSectionEnum.Errors.name] = (
375
+ "No integration test was found in your integration file. A valid Tensorleap "
376
+ "integration file must define an integration test function decorated with "
377
+ "@tensorleap_integration_test. Without it the integration cannot be validated "
378
+ "and the push is aborted. See the mnist leap_integration.py for a valid example."
379
+ )
380
+ return test_result
381
+
359
382
  def _check_preprocess(self) -> DatasetTestResultPayload:
360
383
  test_result = DatasetTestResultPayload('preprocess')
361
384
  try:
@@ -363,7 +386,7 @@ class LeapLoader(LeapLoaderBase):
363
386
  if self.get_sample_id_type() is str:
364
387
  max_allowed_item_size = np.dtype('<U256').itemsize
365
388
  for state, preprocess_response in preprocess_result.items():
366
- sample_ids_array = np.array(preprocess_response.sample_ids)
389
+ sample_ids_array = np.array(preprocess_response.flat_sample_ids)
367
390
  if sample_ids_array.dtype.itemsize > max_allowed_item_size:
368
391
  raise Exception(f"Sample id are too long. Max allowed length is 256 charecters.")
369
392
 
@@ -724,6 +747,57 @@ class LeapLoader(LeapLoaderBase):
724
747
 
725
748
  return sample_ids
726
749
 
750
+ @staticmethod
751
+ def _grouped_dict_keys(rows: List[Dict[str, Any]], handler_name: str) -> List[str]:
752
+ """Keys for a group of dict-metadata rows, requiring every row to share the same keys.
753
+ A group must expose a uniform metadata schema; otherwise the per-key lists would silently
754
+ drop or KeyError on rows that differ from row 0."""
755
+ keys = list(rows[0].keys())
756
+ key_set = set(keys)
757
+ for pos, row in enumerate(rows):
758
+ if set(row.keys()) != key_set:
759
+ raise ValueError(
760
+ f"Metadata handler {handler_name!r} returned inconsistent dict keys across the "
761
+ f"group: row 0 has {sorted(key_set)} but row {pos} has {sorted(row.keys())}. "
762
+ f"All rows in a group must share the same metadata keys.")
763
+ return keys
764
+
765
+ @staticmethod
766
+ def _to_grouped_array(raw: Any) -> Union[npt.NDArray[np.float32], List[npt.NDArray[np.float32]]]:
767
+ """Normalize a grouped encoder result. A list of uniformly-shaped per-sample arrays is
768
+ stacked into a single (B, *) array; a list whose samples differ in shape is kept as a
769
+ list of per-sample arrays (variable-dim groups); an already-grouped array is returned
770
+ as-is."""
771
+ if isinstance(raw, list):
772
+ shapes = {np.asarray(sample).shape for sample in raw}
773
+ if len(shapes) == 1:
774
+ return np.stack(raw, axis=0)
775
+ return raw
776
+ return raw
777
+
778
+ def _locate_group(self, preprocess_state: PreprocessResponse,
779
+ sample_id: Union[int, str]) -> Tuple[List[Union[int, str]], int]:
780
+ """Map a sample id to (its group, its position in that group). Cached per
781
+ PreprocessResponse so grouped single-sample fetches don't rescan the groups."""
782
+ cache = getattr(self, '_group_pos_cache', None)
783
+ if cache is None:
784
+ cache = {}
785
+ self._group_pos_cache = cache
786
+ mapping = cache.get(id(preprocess_state))
787
+ if mapping is None:
788
+ mapping = {}
789
+ for group in preprocess_state.groups:
790
+ for pos, sid in enumerate(group):
791
+ if sid in mapping:
792
+ raise ValueError(
793
+ f"Duplicate sample id {sid!r} across groups in the preprocess response; "
794
+ f"sample ids must be unique within and across groups.")
795
+ mapping[sid] = (group, pos)
796
+ cache[id(preprocess_state)] = mapping
797
+ if sample_id not in mapping:
798
+ raise KeyError(f"Sample id {sample_id!r} is not present in any group of this preprocess response.")
799
+ return mapping[sample_id]
800
+
727
801
  def _get_dataset_handlers(self, handlers: Iterable[DatasetBaseHandler],
728
802
  state: DataStateEnum, sample_id: Union[int, str]) -> Dict[str, npt.NDArray[np.float32]]:
729
803
  result_agg = {}
@@ -735,12 +809,70 @@ class LeapLoader(LeapLoaderBase):
735
809
  sample_id = original_local_id
736
810
  else:
737
811
  preprocess_state = preprocess_result[state]
812
+ if preprocess_state.is_grouped:
813
+ # Grouped dataset: call the encoder once with the whole group, index out the sample.
814
+ group_ids, pos = self._locate_group(preprocess_state, sample_id)
815
+ for handler in handlers:
816
+ grouped = self._to_grouped_array(handler.function(group_ids, preprocess_state))
817
+ result_agg[handler.name] = grouped[pos]
818
+ return result_agg
738
819
  for handler in handlers:
739
820
  handler_result = handler.function(sample_id, preprocess_state)
740
821
  handler_name = handler.name
741
822
  result_agg[handler_name] = handler_result
742
823
  return result_agg
743
824
 
825
+ def get_samples(self, state: DataStateEnum, group_ids: List[Union[int, str]]) -> DatasetSample:
826
+ """Group-aware fetch: hand the whole group to each encoder in a single call and return a
827
+ grouped DatasetSample (inputs/gt as (B, *), metadata as per-key lists of length B, index =
828
+ the group). This is the path a group-aware engine calls to read a file once per group.
829
+ Row order matches group_ids exactly."""
830
+ self.exec_script()
831
+ preprocess_state = self._preprocess_result()[state]
832
+ if preprocess_state.is_grouped:
833
+ # All requested ids must belong to a single group (one file). A cross-group request
834
+ # would force the encoder to load multiple files, defeating the one-load-per-group
835
+ # contract; partitioning a scattered request by group is the engine's responsibility.
836
+ distinct_groups = {id(self._locate_group(preprocess_state, sid)[0]) for sid in group_ids}
837
+ assert len(distinct_groups) == 1, (
838
+ f"get_samples received sample ids spanning {len(distinct_groups)} groups; a request "
839
+ f"must be confined to a single group. The engine partitions cross-group requests and "
840
+ f"issues one get_samples call per group.")
841
+ inputs = {handler.name: self._to_grouped_array(handler.function(group_ids, preprocess_state))
842
+ for handler in global_leap_binder.setup_container.inputs}
843
+ gt = None
844
+ if state != DataStateEnum.unlabeled:
845
+ gt = {handler.name: self._to_grouped_array(handler.function(group_ids, preprocess_state))
846
+ for handler in global_leap_binder.setup_container.ground_truths}
847
+
848
+ metadata: Dict[str, Any] = {}
849
+ metadata_is_none: Dict[str, Any] = {}
850
+ for handler in global_leap_binder.setup_container.metadata:
851
+ rows = handler.function(group_ids, preprocess_state) # list of B scalars/dicts
852
+ if rows and isinstance(rows[0], dict):
853
+ for key in self._grouped_dict_keys(rows, handler.name):
854
+ name = "{}_{}".format(handler.name, key)
855
+ converted = [self._convert_metadata_to_correct_type(name, row[key]) for row in rows]
856
+ metadata[name] = [v for v, _ in converted]
857
+ metadata_is_none[name] = [is_none for _, is_none in converted]
858
+ else:
859
+ converted = [self._convert_metadata_to_correct_type(handler.name, row) for row in rows]
860
+ metadata[handler.name] = [v for v, _ in converted]
861
+ metadata_is_none[handler.name] = [is_none for _, is_none in converted]
862
+
863
+ # custom_latent_space is group-aware like the input/GT encoders: hand it the whole group in a
864
+ # single call so a file-backed latent fn loads the file once, and normalize to (B, d). This
865
+ # keeps the mandatory group path non-lossy. instance_masks stay None here: they are a
866
+ # per-(sample, instance) concern that needs an instance_id the group fetch does not carry.
867
+ custom_latent_space = None
868
+ if global_leap_binder.setup_container.custom_latent_space is not None:
869
+ latent_fn = global_leap_binder.setup_container.custom_latent_space.function
870
+ custom_latent_space = self._to_grouped_array(latent_fn(group_ids, preprocess_state))
871
+
872
+ return DatasetSample(inputs=inputs, gt=gt, metadata=metadata, metadata_is_none=metadata_is_none,
873
+ index=list(group_ids), state=state, custom_latent_space=custom_latent_space,
874
+ instance_masks=None)
875
+
744
876
  def _get_inputs(self, state: DataStateEnum, sample_id: Union[int, str]) -> Dict[str, npt.NDArray[np.float32]]:
745
877
  return self._get_dataset_handlers(global_leap_binder.setup_container.inputs, state, sample_id)
746
878
 
@@ -821,12 +953,18 @@ class LeapLoader(LeapLoaderBase):
821
953
  sample_id = original_local_id
822
954
  else:
823
955
  preprocess_state = preprocess_result[state]
956
+ group_ids, pos = (None, None)
957
+ if preprocess_state.is_grouped:
958
+ group_ids, pos = self._locate_group(preprocess_state, sample_id)
824
959
  for handler in global_leap_binder.setup_container.metadata:
825
960
  if requested_metadata_names:
826
961
  if not is_metadata_name_starts_with_handler_name(handler):
827
962
  continue
828
963
 
829
- handler_result = handler.function(sample_id, preprocess_state)
964
+ if preprocess_state.is_grouped:
965
+ handler_result = handler.function(group_ids, preprocess_state)[pos]
966
+ else:
967
+ handler_result = handler.function(sample_id, preprocess_state)
830
968
  if isinstance(handler_result, dict):
831
969
  for single_metadata_name, single_metadata_result in handler_result.items():
832
970
  handler_name = f'{handler.name}_{single_metadata_name}'
@@ -845,6 +983,66 @@ class LeapLoader(LeapLoaderBase):
845
983
 
846
984
  return result_agg, is_none
847
985
 
986
+ def _grouped_metadata_for_group(self, preprocess_state: PreprocessResponse,
987
+ group_ids: List[Union[int, str]],
988
+ requested_metadata_names: Optional[List[str]]
989
+ ) -> Tuple[Dict[str, List[Any]], Dict[str, List[bool]]]:
990
+ def is_wanted(_handler):
991
+ if not requested_metadata_names:
992
+ return True
993
+ for metadata_name in requested_metadata_names:
994
+ if metadata_name.startswith(_handler.name + '_') or metadata_name == _handler.name:
995
+ return True
996
+ return False
997
+
998
+ results: Dict[str, List[Any]] = {}
999
+ is_none: Dict[str, List[bool]] = {}
1000
+ for handler in global_leap_binder.setup_container.metadata:
1001
+ if requested_metadata_names and not is_wanted(handler):
1002
+ continue
1003
+ rows = handler.function(group_ids, preprocess_state)
1004
+ if rows and isinstance(rows[0], dict):
1005
+ for key in self._grouped_dict_keys(rows, handler.name):
1006
+ name = f'{handler.name}_{key}'
1007
+ if requested_metadata_names and name not in requested_metadata_names:
1008
+ continue
1009
+ converted = [self._convert_metadata_to_correct_type(name, row[key]) for row in rows]
1010
+ results[name] = [v for v, _ in converted]
1011
+ is_none[name] = [n for _, n in converted]
1012
+ else:
1013
+ name = handler.name
1014
+ if requested_metadata_names and name not in requested_metadata_names:
1015
+ continue
1016
+ converted = [self._convert_metadata_to_correct_type(name, row) for row in rows]
1017
+ results[name] = [v for v, _ in converted]
1018
+ is_none[name] = [n for _, n in converted]
1019
+ return results, is_none
1020
+
1021
+ def get_metadata_multiple_samples(self, state: DataStateEnum, sample_ids: Union[List[int], List[str]],
1022
+ requested_metadata_names: Optional[List[str]] = None
1023
+ ) -> Tuple[Dict[str, Union[List[str], List[int], List[bool], List[float]]],
1024
+ Dict[str, List[bool]]]:
1025
+ preprocess_state = self._preprocess_result().get(state)
1026
+ if preprocess_state is None or not preprocess_state.is_grouped:
1027
+ return super().get_metadata_multiple_samples(state, sample_ids, requested_metadata_names)
1028
+
1029
+ sample_id_type = self.get_sample_id_type()
1030
+ aggregated_results: Dict[str, List[Any]] = {}
1031
+ aggregated_is_none: Dict[str, List[bool]] = {}
1032
+ group_cache: Dict[int, Tuple[Dict[str, List[Any]], Dict[str, List[bool]]]] = {}
1033
+ for sample_id in sample_ids:
1034
+ sample_id = sample_id_type(sample_id)
1035
+ group_ids, pos = self._locate_group(preprocess_state, sample_id)
1036
+ key = id(group_ids)
1037
+ if key not in group_cache:
1038
+ group_cache[key] = self._grouped_metadata_for_group(
1039
+ preprocess_state, group_ids, requested_metadata_names)
1040
+ group_results, group_is_none = group_cache[key]
1041
+ for name, values in group_results.items():
1042
+ aggregated_results.setdefault(name, []).append(values[pos])
1043
+ aggregated_is_none.setdefault(name, []).append(group_is_none[name][pos])
1044
+ return aggregated_results, aggregated_is_none
1045
+
848
1046
  @lru_cache()
849
1047
  def get_sample_id_type(self) -> Type:
850
1048
  preprocess_results = list(self._preprocess_result().values())
code_loader/utils.py CHANGED
@@ -14,6 +14,11 @@ from code_loader.contract.enums import DatasetMetadataType
14
14
  def to_numpy_return_wrapper(encoder_function: SectionCallableInterface) -> SectionCallableInterface:
15
15
  def numpy_encoder_function(idx: Union[int, str], samples: PreprocessResponse) -> npt.NDArray[np.float32]:
16
16
  result = encoder_function(idx, samples)
17
+ if isinstance(result, list) and len(result) > 0 and all(isinstance(sample, np.ndarray) for sample in result):
18
+ if len({sample.shape for sample in result}) > 1:
19
+ # Variable per-sample shapes within a group: keep the per-sample arrays as a list
20
+ # instead of forcing a stack that would raise an inhomogeneous-shape error.
21
+ return result
17
22
  numpy_result: npt.NDArray[np.float32] = np.array(result)
18
23
  return numpy_result
19
24
 
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2021 TensorLeap
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -1,7 +1,8 @@
1
- Metadata-Version: 2.3
1
+ Metadata-Version: 2.1
2
2
  Name: code-loader
3
- Version: 1.0.190.dev0
3
+ Version: 1.0.191.dev1
4
4
  Summary:
5
+ Home-page: https://github.com/tensorleap/code-loader
5
6
  License: MIT
6
7
  Author: dorhar
7
8
  Author-email: doron.harnoy@tensorleap.ai
@@ -19,7 +20,6 @@ Requires-Dist: numpy (>=2.3.2,<3.0.0) ; python_version >= "3.11" and python_vers
19
20
  Requires-Dist: psutil (>=5.9.5,<6.0.0)
20
21
  Requires-Dist: pyyaml (>=6.0.2,<7.0.0)
21
22
  Requires-Dist: requests (>=2.32.3,<3.0.0)
22
- Project-URL: Homepage, https://github.com/tensorleap/code-loader
23
23
  Project-URL: Repository, https://github.com/tensorleap/code-loader
24
24
  Description-Content-Type: text/markdown
25
25
 
@@ -1,6 +1,7 @@
1
+ LICENSE,sha256=qIwWjdspQeSMTtnFZBC8MuT-95L02FPvzRUdWFxrwJY,1067
1
2
  code_loader/__init__.py,sha256=outxRQ0M-zMfV0QGVJmAed5qWfRmyD0TV6-goEGAzBw,406
2
3
  code_loader/contract/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
3
- code_loader/contract/datasetclasses.py,sha256=RYZcnX7vLMYhyQtCbX2PscVFNAYULxgAA9U4DRRfLqA,10456
4
+ code_loader/contract/datasetclasses.py,sha256=xkem0ONbgsTWEG_4G-hdkFLJie-xSXSQZBzsZcSjQbU,12724
4
5
  code_loader/contract/enums.py,sha256=2q-IV_5g9lLE306DIbWA1c0tn5IhDtxsKxyV1x_Lreg,1671
5
6
  code_loader/contract/exceptions.py,sha256=jWqu5i7t-0IG0jGRsKF4DjJdrsdpJjIYpUkN1F4RiyQ,51
6
7
  code_loader/contract/mapping.py,sha256=sWJhpng-IkOzQnWQdMT5w2ZZ3X1Z_OOzSwCLXIS7oxE,1446
@@ -20,18 +21,18 @@ code_loader/experiment_api/types.py,sha256=MY8xFARHwdVA7p4dxyhD60ShmttgTvb4qdp1o
20
21
  code_loader/experiment_api/utils.py,sha256=XZHtxge12TS4H4-8PjV3sKuhp8Ud6ojAiIzTZJEqBqc,3304
21
22
  code_loader/experiment_api/workingspace_config_utils.py,sha256=DLzXQCg4dgTV_YgaSbeTVzq-2ja_SQw4zi7LXwKL9cY,990
22
23
  code_loader/inner_leap_binder/__init__.py,sha256=koOlJyMNYzGbEsoIbXathSmQ-L38N_pEXH_HvL7beXU,99
23
- code_loader/inner_leap_binder/leapbinder.py,sha256=PAX9cmT1QQV7PWsbY9rMRWktpM7KryfFlwBiSwSskzg,42654
24
- code_loader/inner_leap_binder/leapbinder_decorators.py,sha256=lgKJ9cJDuDABscDqLvnqRt7ViU6YeZFcm9996EUosGA,123107
25
- code_loader/leaploader.py,sha256=K9Q5kiCZ_A2GkS5qOEantAMvWGaz3bmO1X8buXDSpgg,44682
24
+ code_loader/inner_leap_binder/leapbinder.py,sha256=pds20IP5s1iBXUIGbweDxiq77x4DIRjNu110a7Xgpc4,43590
25
+ code_loader/inner_leap_binder/leapbinder_decorators.py,sha256=w29kFpQydGtdBE8joHW0lCKhsMF91PyeodoXgqsR_rw,128103
26
+ code_loader/leaploader.py,sha256=WhaV4TI2fRPdo72NrK0y3NaaXStBaKviGTsoSADyom4,56652
26
27
  code_loader/leaploaderbase.py,sha256=l36qDA00GhZEG5NLKpEtAXgWJA-UQQIhNFGxywK7mUA,6530
27
28
  code_loader/mixpanel_tracker.py,sha256=rNwRmFifNbdUoqLQvvhhgpKczWpWiEmd8MfyJe27sxw,9131
28
29
  code_loader/plot_functions/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
29
30
  code_loader/plot_functions/plot_functions.py,sha256=2DC-zlVaN13P4VNx5d8csgs80C6SisaeP1-Kq2LW7iM,16075
30
31
  code_loader/plot_functions/visualize.py,sha256=gsBAYYkwMh7jIpJeDMPS8G4CW-pxwx6LznoQIvi4vpo,657
31
- code_loader/utils.py,sha256=YecipkdTA-VcE9F0RQcY9cFnY8P3AksPnHM2Db7xUSk,3972
32
+ code_loader/utils.py,sha256=PXRsmzFt_iPQIOQ2DyuyNnIdFW0tJElE5GsTnwY4Enk,4371
32
33
  code_loader/visualizers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
33
34
  code_loader/visualizers/default_visualizers.py,sha256=onRnLE_TXfgLN4o52hQIOOhUcFexGlqJ3xSpQDVLuZM,2604
34
- code_loader-1.0.190.dev0.dist-info/LICENSE,sha256=qIwWjdspQeSMTtnFZBC8MuT-95L02FPvzRUdWFxrwJY,1067
35
- code_loader-1.0.190.dev0.dist-info/METADATA,sha256=2Tqb-fK_hBKdjefZN31fYytgQdg3VS9rgcPiwqLsLkc,1107
36
- code_loader-1.0.190.dev0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
37
- code_loader-1.0.190.dev0.dist-info/RECORD,,
35
+ code_loader-1.0.191.dev1.dist-info/LICENSE,sha256=qIwWjdspQeSMTtnFZBC8MuT-95L02FPvzRUdWFxrwJY,1067
36
+ code_loader-1.0.191.dev1.dist-info/METADATA,sha256=MM9kKQMqR6Uo7iZ8ca6ZYNAZmUGtGxvTr6cjWwEhlkU,1095
37
+ code_loader-1.0.191.dev1.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
38
+ code_loader-1.0.191.dev1.dist-info/RECORD,,
@@ -1,4 +1,4 @@
1
1
  Wheel-Version: 1.0
2
- Generator: poetry-core 2.1.3
2
+ Generator: poetry-core 1.9.0
3
3
  Root-Is-Purelib: true
4
4
  Tag: py3-none-any
File without changes