Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / data / indexed / indexed_simple_blended_dataset.py: 0%

1# Copyright 2026 Huawei Technologies Co., Ltd 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14# ============================================================================ 

15"""Index-free blending for pre-shuffled MR Datasets.""" 

16 

17from __future__ import annotations 

18 

19from collections.abc import Mapping, Sequence 

20from typing import Any, Literal 

21 

22SimpleBlendMode = Literal["inter", "intra"] 

23 

24 

25class SimpleBlendedDataset: 

26 """Blend equally weighted MR Datasets in interleaved or contiguous order.""" 

27 

28 def __init__( 

29 self, 

30 datasets: Sequence[Any], 

31 size: int, 

32 mode: SimpleBlendMode, 

33 ) -> None: 

34 """Validate the inputs and create the configured simple ordering.""" 

35 if not datasets: 

36 raise ValueError("simple blend requires at least one Dataset") 

37 self.datasets = list(datasets) 

38 self.size = size 

39 self.mode = mode 

40 self._locations = self._build_locations() 

41 if self.size > len(self._locations): 

42 raise ValueError( 

43 f"Requested {self.size} samples from a simple blend containing " 

44 f"{len(self._locations)} samples" 

45 ) 

46 

47 def __len__(self) -> int: 

48 """Return the requested blended sample count.""" 

49 return self.size 

50 

51 def __getitem__(self, index: int) -> Mapping[str, Any]: 

52 """Read one MR sample in interleaved or Dataset-contiguous order.""" 

53 dataset_id, sample_id = self._locations[index] 

54 sample = {"dataset_id": dataset_id, **self.datasets[dataset_id][sample_id]} 

55 return sample 

56 

57 def _build_locations(self) -> list[tuple[int, int]]: 

58 """Map blend positions to Dataset and sample indices.""" 

59 if self.mode == "intra": 

60 locations = [ 

61 (dataset_id, sample_id) 

62 for dataset_id, dataset in enumerate(self.datasets) 

63 for sample_id in range(len(dataset)) 

64 ] 

65 return locations 

66 

67 # mode = inter 

68 locations = self._build_interleaved_locations() 

69 return locations 

70 

71 def _build_interleaved_locations(self) -> list[tuple[int, int]]: 

72 """Alternate over non-empty Datasets until every sample is exposed.""" 

73 max_size = max(len(dataset) for dataset in self.datasets) 

74 locations = [] 

75 for sample_id in range(max_size): 

76 for dataset_id, dataset in enumerate(self.datasets): 

77 if sample_id < len(dataset): 

78 locations.append((dataset_id, sample_id)) 

79 return locations