code-loader 1.0.191__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):
@@ -26,12 +26,13 @@ from code_loader.contract.enums import MetricDirection, LeapDataType, DatasetMet
26
26
  from code_loader import leap_binder, LeapLoader
27
27
  from code_loader.contract.mapping import NodeMapping, NodeMappingType, NodeConnection
28
28
  from code_loader.contract.visualizer_classes import LeapImage, LeapImageMask, LeapTextMask, LeapText, LeapGraph, \
29
- LeapHorizontalBar, LeapImageWithBBox, LeapImageWithHeatmap, LeapVideo, LeapValidationError
29
+ LeapHorizontalBar, LeapImageWithBBox, LeapImageWithHeatmap, LeapVideo
30
30
  from code_loader.inner_leap_binder.leapbinder import mapping_runtime_mode_env_var_mame
31
31
  from code_loader.mixpanel_tracker import clear_integration_events, AnalyticsEvent, emit_integration_event_once
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))
@@ -316,17 +359,18 @@ def tensorleap_integration_test():
316
359
  first_tb = traceback.extract_tb(e.__traceback__)[-1]
317
360
  file_name = Path(first_tb.filename).name
318
361
  line_number = first_tb.lineno
319
- update_env_params_func("code_mapping", "x")
320
- if isinstance(e, LeapValidationError):
321
- raise Exception(f'Invalid integration code. File {file_name}, line {line_number}: {e}')
322
- elif isinstance(e, TypeError) and 'is not subscriptable' in str(e):
362
+ if isinstance(e, TypeError) and 'is not subscriptable' in str(e):
363
+ update_env_params_func("code_mapping", "x")
323
364
  raise Exception(f'Invalid integration code. File {file_name}, line {line_number}: '
324
- 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.")
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.")
325
366
  else:
367
+ update_env_params_func("code_mapping", "x")
368
+
326
369
  raise Exception(f'Invalid integration code. File {file_name}, line {line_number}: '
327
- f'Integration test is only allowed to call Tensorleap decorators. '
328
- f'Ensure any arithmetics, external library use, Python logic is placed within Tensorleap decoders')
370
+ f'Integration test is only allowed to call Tensorleap decorators. '
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:
@@ -611,9 +647,7 @@ def tensorleap_load_model(prediction_types: Optional[List[PredictionTypeHandler]
611
647
 
612
648
  # onnx runtime interface
613
649
  def run(self, output_names, input_dict):
614
- if output_names is not None:
615
- raise LeapValidationError("Tensorleap doesn't support selecting outputs by name — "
616
- "use model.run(None, inputs) to return all outputs.")
650
+ assert output_names is None
617
651
  assert isinstance(input_dict, dict), \
618
652
  f'Expected input_dict to be a dict, got {type(input_dict)} instead.'
619
653
  for i, (input_key, elem) in enumerate(input_dict.items()):
@@ -1518,13 +1552,10 @@ def tensorleap_metadata(
1518
1552
  raise Exception(f'Metadata with name {name} already exists. '
1519
1553
  f'Please choose another')
1520
1554
 
1521
- def _validate_input_args(sample_id: Union[int, str], preprocess_response: PreprocessResponse):
1522
- assert type(sample_id) == preprocess_response.sample_id_type, \
1523
- (f'{user_function.__name__}() validation failed: '
1524
- f'Argument sample_id should be as the same type as defined in the preprocess response '
1525
- 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__)
1526
1557
 
1527
- def _validate_result(result):
1558
+ def _validate_single(result):
1528
1559
  supported_result_types = (type(None), int, str, bool, float, dict, np.floating,
1529
1560
  np.bool_, np.unsignedinteger, np.signedinteger, np.integer)
1530
1561
  if isinstance(result, tuple):
@@ -1542,6 +1573,20 @@ def tensorleap_metadata(
1542
1573
  (f'{user_function.__name__}() validation failed: '
1543
1574
  f'Values in the return dict should be of type {str(supported_result_types)}. Got {type(value)}.')
1544
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
+
1545
1590
  def inner_without_validate(sample_id, preprocess_response):
1546
1591
 
1547
1592
  global _called_from_inside_tl_decorator
@@ -1561,14 +1606,16 @@ def tensorleap_metadata(
1561
1606
  set_current('tensorleap_metadata')
1562
1607
  if os.environ.get(mapping_runtime_mode_env_var_mame):
1563
1608
  return None
1564
- validate_args_structure(*args, types_order=[Union[int, str], PreprocessResponse],
1609
+ validate_args_structure(*args, types_order=[Union[int, str, list], PreprocessResponse],
1565
1610
  func_name=user_function.__name__, expected_names=["idx", "preprocess"], **kwargs)
1566
1611
  sample_id, preprocess_response = args if len(args) != 0 else kwargs.values()
1567
1612
  _validate_input_args(sample_id, preprocess_response)
1568
1613
 
1614
+ grouped = isinstance(sample_id, list)
1615
+ group_size = len(sample_id) if grouped else None
1569
1616
  result = inner_without_validate(sample_id, preprocess_response)
1570
1617
 
1571
- _validate_result(result)
1618
+ _validate_result(result, grouped, group_size)
1572
1619
  if not _call_from_tl_platform:
1573
1620
  update_env_params_func("tensorleap_metadata", "v")
1574
1621
  return result
@@ -1580,35 +1627,50 @@ def tensorleap_metadata(
1580
1627
 
1581
1628
  def tensorleap_custom_latent_space():
1582
1629
  def decorating_function(user_function: SectionCallableInterface):
1583
- def _validate_input_args(sample_id: Union[int, str], preprocess_response: PreprocessResponse):
1584
- assert isinstance(sample_id, (int, str)), \
1585
- (f'tensorleap_custom_latent_space validation failed: '
1586
- f'Argument sample_id should be either int or str. Got {type(sample_id)}.')
1587
- assert isinstance(preprocess_response, PreprocessResponse), \
1588
- (f'tensorleap_custom_latent_space validation failed: '
1589
- f'Argument preprocess_response should be a PreprocessResponse. Got {type(preprocess_response)}.')
1590
- assert type(sample_id) == preprocess_response.sample_id_type, \
1591
- (f'tensorleap_custom_latent_space validation failed: '
1592
- f'Argument sample_id should be as the same type as defined in the preprocess response '
1593
- 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')
1594
1632
 
1595
- def _validate_result(result):
1596
- assert isinstance(result, np.ndarray), \
1633
+ def _validate_single(single_result):
1634
+ assert isinstance(single_result, np.ndarray), \
1597
1635
  (f'tensorleap_custom_latent_space validation failed: '
1598
- f'The return type should be a numpy array. Got {type(result)}.')
1599
- if result.ndim > 1:
1600
- 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))
1601
1639
  store_general_warning(
1602
- key=("tensorleap_custom_latent_space_flatten", tuple(result.shape)),
1640
+ key=("tensorleap_custom_latent_space_flatten", tuple(single_result.shape)),
1603
1641
  message=(
1604
- f"tensorleap_custom_latent_space returned per-sample shape {tuple(result.shape)} "
1605
- 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 "
1606
1644
  f"flatten to ({flat_dim},) before downstream visualization and clustering. "
1607
1645
  f"If you want a different aggregation (e.g. global average pooling), do it "
1608
1646
  f"inside your function."
1609
1647
  ),
1610
1648
  )
1611
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
+
1612
1674
  def inner_without_validate(sample_id, preprocess_response):
1613
1675
  global _called_from_inside_tl_decorator
1614
1676
  _called_from_inside_tl_decorator += 1
@@ -1628,9 +1690,11 @@ def tensorleap_custom_latent_space():
1628
1690
 
1629
1691
  _validate_input_args(sample_id, preprocess_response)
1630
1692
 
1693
+ grouped = isinstance(sample_id, list)
1694
+ group_size = len(sample_id) if grouped else None
1631
1695
  result = inner_without_validate(sample_id, preprocess_response)
1632
1696
 
1633
- _validate_result(result)
1697
+ _validate_result(result, grouped, group_size)
1634
1698
  return result
1635
1699
 
1636
1700
  return inner
@@ -1668,6 +1732,15 @@ def tensorleap_preprocess():
1668
1732
  assert len(set(result)) == len(result), \
1669
1733
  (f'{user_function.__name__}() validation failed: '
1670
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}.')
1671
1744
 
1672
1745
  def inner(*args, **kwargs):
1673
1746
  if not _call_from_tl_platform:
@@ -2005,14 +2078,10 @@ def tensorleap_input_encoder(name: str, channel_dim=_UNSET, model_input_index=No
2005
2078
  f"channel axis in the batched tensor (batch = axis 0), so 0 is invalid. Use 1 "
2006
2079
  f"for channel-first (NCHW) and -1 for channel-last (NHWC).")
2007
2080
 
2008
- def _validate_input_args(sample_id: Union[int, str], preprocess_response: PreprocessResponse):
2009
- assert type(sample_id) == preprocess_response.sample_id_type, \
2010
- (f'{user_function.__name__}() validation failed: '
2011
- f'Argument sample_id should be as the same type as defined in the preprocess response '
2012
- 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__)
2013
2083
 
2014
- def _validate_result(result):
2015
- validate_output_structure(result, func_name=user_function.__name__, expected_type_name="np.ndarray")
2084
+ def _validate_single(result):
2016
2085
  assert isinstance(result, np.ndarray), \
2017
2086
  (f'{user_function.__name__}() validation failed: '
2018
2087
  f'Unsupported return type. Should be a numpy array. Got {type(result)}.')
@@ -2022,6 +2091,13 @@ def tensorleap_input_encoder(name: str, channel_dim=_UNSET, model_input_index=No
2022
2091
  assert channel_dim - 1 <= len(result.shape), (f'{user_function.__name__}() validation failed: '
2023
2092
  f'The channel_dim ({channel_dim}) should be <= to the rank of the resulting input rank ({len(result.shape)}).')
2024
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
+
2025
2101
  def inner_without_validate(sample_id, preprocess_response):
2026
2102
  global _called_from_inside_tl_decorator
2027
2103
  _called_from_inside_tl_decorator += 1
@@ -2038,16 +2114,18 @@ def tensorleap_input_encoder(name: str, channel_dim=_UNSET, model_input_index=No
2038
2114
  def inner(*args, **kwargs):
2039
2115
  if not _call_from_tl_platform:
2040
2116
  set_current("tensorleap_input_encoder")
2041
- validate_args_structure(*args, types_order=[Union[int, str], PreprocessResponse],
2117
+ validate_args_structure(*args, types_order=[Union[int, str, list], PreprocessResponse],
2042
2118
  func_name=user_function.__name__, expected_names=["idx", "preprocess"], **kwargs)
2043
2119
  sample_id, preprocess_response = args if len(args) != 0 else kwargs.values()
2044
2120
  _validate_input_args(sample_id, preprocess_response)
2045
2121
 
2122
+ grouped = isinstance(sample_id, list)
2123
+ group_size = len(sample_id) if grouped else None
2046
2124
  result = inner_without_validate(sample_id, preprocess_response)
2047
2125
 
2048
- _validate_result(result)
2126
+ _validate_result(result, grouped, group_size)
2049
2127
 
2050
- 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:
2051
2129
  batch_warning(result, user_function.__name__)
2052
2130
  result = np.expand_dims(result, axis=0)
2053
2131
  # Emit integration test event once per test
@@ -2070,8 +2148,15 @@ def tensorleap_input_encoder(name: str, channel_dim=_UNSET, model_input_index=No
2070
2148
  inner.node_mapping = NodeMapping(name, node_mapping_type)
2071
2149
 
2072
2150
  def mapping_inner(*args, **kwargs):
2073
- class TempMapping:
2074
- 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
2075
2160
 
2076
2161
  ret = TempMapping()
2077
2162
  ret.node_mapping = mapping_inner.node_mapping
@@ -2101,15 +2186,10 @@ def tensorleap_gt_encoder(name: str):
2101
2186
  raise Exception(f'GT with name {name} already exists. '
2102
2187
  f'Please choose another')
2103
2188
 
2104
- def _validate_input_args(sample_id: Union[int, str], preprocess_response: PreprocessResponse):
2105
- assert type(sample_id) == preprocess_response.sample_id_type, \
2106
- (f'{user_function.__name__}() validation failed: '
2107
- f'Argument sample_id should be as the same type as defined in the preprocess response '
2108
- 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__)
2109
2191
 
2110
- def _validate_result(result):
2111
- validate_output_structure(result, func_name=user_function.__name__, expected_type_name="np.ndarray",
2112
- gt_flag=True)
2192
+ def _validate_single(result):
2113
2193
  assert isinstance(result, np.ndarray), \
2114
2194
  (f'{user_function.__name__}() validation failed: '
2115
2195
  f'Unsupported return type. Should be a numpy array. Got {type(result)}.')
@@ -2117,6 +2197,14 @@ def tensorleap_gt_encoder(name: str):
2117
2197
  (f'{user_function.__name__}() validation failed: '
2118
2198
  f'The return type should be a numpy array of type float32. Got {result.dtype}.')
2119
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
+
2120
2208
  def inner_without_validate(sample_id, preprocess_response):
2121
2209
  global _called_from_inside_tl_decorator
2122
2210
  _called_from_inside_tl_decorator += 1
@@ -2133,16 +2221,18 @@ def tensorleap_gt_encoder(name: str):
2133
2221
  def inner(*args, **kwargs):
2134
2222
  if not _call_from_tl_platform:
2135
2223
  set_current("tensorleap_gt_encoder")
2136
- validate_args_structure(*args, types_order=[Union[int, str], PreprocessResponse],
2224
+ validate_args_structure(*args, types_order=[Union[int, str, list], PreprocessResponse],
2137
2225
  func_name=user_function.__name__, expected_names=["idx", "preprocess"], **kwargs)
2138
2226
  sample_id, preprocess_response = args
2139
2227
  _validate_input_args(sample_id, preprocess_response)
2140
2228
 
2229
+ grouped = isinstance(sample_id, list)
2230
+ group_size = len(sample_id) if grouped else None
2141
2231
  result = inner_without_validate(sample_id, preprocess_response)
2142
2232
 
2143
- _validate_result(result)
2233
+ _validate_result(result, grouped, group_size)
2144
2234
 
2145
- 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:
2146
2236
  batch_warning(result, user_function.__name__)
2147
2237
  result = np.expand_dims(result, axis=0)
2148
2238
  # Emit integration test event once per test
@@ -2159,8 +2249,15 @@ def tensorleap_gt_encoder(name: str):
2159
2249
  inner.node_mapping = NodeMapping(name, NodeMappingType.GroundTruth)
2160
2250
 
2161
2251
  def mapping_inner(*args, **kwargs):
2162
- class TempMapping:
2163
- 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
2164
2261
 
2165
2262
  ret = TempMapping()
2166
2263
  ret.node_mapping = mapping_inner.node_mapping
@@ -2437,14 +2534,14 @@ def tensorleap_status_table():
2437
2534
  ready = False
2438
2535
  next_step = row["name"]
2439
2536
 
2440
- if code_mapping_failure[0]:
2441
- print(f"\n{CROSS} {code_mapping_failure_mes}")
2442
- return
2443
-
2444
2537
  if _crashed["value"]:
2445
2538
  print(f"\nScript crashed before completing all steps. crashed at function '{_current_func['name']}'.")
2446
2539
  return
2447
2540
 
2541
+ if code_mapping_failure[0]:
2542
+ print(f"\n{CROSS + code_mapping_failure_mes}.")
2543
+ return
2544
+
2448
2545
  print(ready_mess) if ready else print(
2449
2546
  mandatory_ready_mess + next_step
2450
2547
  ) if (next_step and "optional" in next_step) else print(not_ready_mess + (next_step or ""))
@@ -2457,10 +2554,6 @@ def tensorleap_status_table():
2457
2554
  code_mapping_failure[0] = 1
2458
2555
  if status == "v":
2459
2556
  _set_status(name, CHECK)
2460
- # A decorator that reports success is no longer the "currently running"
2461
- # one, so clear the pointer. This keeps a later body-level raise from
2462
- # being blamed on the last decorator that actually succeeded.
2463
- _current_func["name"] = None
2464
2557
  else:
2465
2558
  _set_status(name, CROSS)
2466
2559
 
@@ -2507,14 +2600,10 @@ def tensorleap_status_table():
2507
2600
  def handle_exception(exc_type, exc_value, exc_traceback):
2508
2601
  _crashed["value"] = True
2509
2602
  crashed_name = _current_func["name"]
2510
- if crashed_name and not code_mapping_failure[0]:
2603
+ if crashed_name:
2511
2604
  row = _find_row(crashed_name)
2512
2605
  if row:
2513
2606
  row["Added to integration"] = CROSS
2514
- elif not crashed_name:
2515
- # Crash with no decorator currently running = the integration-test body /
2516
- # code flow itself failed; report that rather than blaming a decorator.
2517
- code_mapping_failure[0] = 1
2518
2607
 
2519
2608
  traceback.print_exception(exc_type, exc_value, exc_traceback)
2520
2609
  run_on_exit()
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),
@@ -382,7 +386,7 @@ class LeapLoader(LeapLoaderBase):
382
386
  if self.get_sample_id_type() is str:
383
387
  max_allowed_item_size = np.dtype('<U256').itemsize
384
388
  for state, preprocess_response in preprocess_result.items():
385
- sample_ids_array = np.array(preprocess_response.sample_ids)
389
+ sample_ids_array = np.array(preprocess_response.flat_sample_ids)
386
390
  if sample_ids_array.dtype.itemsize > max_allowed_item_size:
387
391
  raise Exception(f"Sample id are too long. Max allowed length is 256 charecters.")
388
392
 
@@ -743,6 +747,57 @@ class LeapLoader(LeapLoaderBase):
743
747
 
744
748
  return sample_ids
745
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
+
746
801
  def _get_dataset_handlers(self, handlers: Iterable[DatasetBaseHandler],
747
802
  state: DataStateEnum, sample_id: Union[int, str]) -> Dict[str, npt.NDArray[np.float32]]:
748
803
  result_agg = {}
@@ -754,12 +809,70 @@ class LeapLoader(LeapLoaderBase):
754
809
  sample_id = original_local_id
755
810
  else:
756
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
757
819
  for handler in handlers:
758
820
  handler_result = handler.function(sample_id, preprocess_state)
759
821
  handler_name = handler.name
760
822
  result_agg[handler_name] = handler_result
761
823
  return result_agg
762
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
+
763
876
  def _get_inputs(self, state: DataStateEnum, sample_id: Union[int, str]) -> Dict[str, npt.NDArray[np.float32]]:
764
877
  return self._get_dataset_handlers(global_leap_binder.setup_container.inputs, state, sample_id)
765
878
 
@@ -840,12 +953,18 @@ class LeapLoader(LeapLoaderBase):
840
953
  sample_id = original_local_id
841
954
  else:
842
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)
843
959
  for handler in global_leap_binder.setup_container.metadata:
844
960
  if requested_metadata_names:
845
961
  if not is_metadata_name_starts_with_handler_name(handler):
846
962
  continue
847
963
 
848
- 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)
849
968
  if isinstance(handler_result, dict):
850
969
  for single_metadata_name, single_metadata_result in handler_result.items():
851
970
  handler_name = f'{handler.name}_{single_metadata_name}'
@@ -864,6 +983,66 @@ class LeapLoader(LeapLoaderBase):
864
983
 
865
984
  return result_agg, is_none
866
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
+
867
1046
  @lru_cache()
868
1047
  def get_sample_id_type(self) -> Type:
869
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.191
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=yMEROthjXugzqJCCjL5m7Erkv33L56fqACDXaFK2M4c,123503
25
- code_loader/leaploader.py,sha256=9bUl69X5pk6prgFMCnaoA4MhRoacDb9Ddv3CWgG1F0Y,45968
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.191.dist-info/LICENSE,sha256=qIwWjdspQeSMTtnFZBC8MuT-95L02FPvzRUdWFxrwJY,1067
35
- code_loader-1.0.191.dist-info/METADATA,sha256=HPzc6iqSgc1a5d-yKe2kUtYbcnqcpVMKfMQxiS9Rb14,1102
36
- code_loader-1.0.191.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
37
- code_loader-1.0.191.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