Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / trainer / config / data.py: 96%

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"""Model-assets, dataset and dataloader configuration sections. 

16 

17Split from ``auto_models/trainer/config.py`` in stage 7 (05 §15.2.5); 

18class names, fields and defaults are unchanged. 

19""" 

20 

21from dataclasses import dataclass, field 

22from typing import Any, Literal, Optional, Union 

23 

24from hyper_parallel.trainer.config.target import Target, _serialize_config_value 

25 

26 

27@dataclass 

28class ModelAssetsConfig: 

29 """Tokenizer and chat-template configuration for text datasets.""" 

30 

31 chat_template: Optional[Union[str, Target[Any]]] = None 

32 chat_template_kwargs: dict[str, Any] = field(default_factory=dict) 

33 tokenizer: Optional[Target[Any]] = None 

34 

35 

36@dataclass 

37class DatasetConfig: 

38 """Dataset target with its model assets and sample transform.""" 

39 

40 target: Target[Any] 

41 model_assets: Union[ModelAssetsConfig, Target[Any]] = field(default_factory=ModelAssetsConfig) 

42 data_transform: Optional[Target[Any]] = None 

43 training_mode: Optional[Literal["sft", "pretrain"]] = None 

44 

45 def build(self, **runtime_kwargs: Any) -> Any: 

46 """Build the Dataset target with runtime Trainer arguments.""" 

47 dataset = self.target.build(**runtime_kwargs) 

48 return dataset 

49 

50 def __getattr__(self, name: str) -> Any: 

51 """Expose configured Dataset options through the wrapped target.""" 

52 return getattr(self.target, name) 

53 

54 def to_dict(self) -> dict[str, Any]: 

55 """Serialize Dataset components in their compact nested YAML shape.""" 

56 config = self.target.to_dict() 

57 config["model_assets"] = _serialize_config_value(self.model_assets) 

58 config["data_transform"] = _serialize_config_value(self.data_transform) 

59 if self.training_mode is not None: 

60 config["training_mode"] = self.training_mode 

61 return config 

62 

63 

64@dataclass 

65class DataLoaderConfig: 

66 """DataLoader target and its text-batch assembly components. 

67 

68 ``load_balance=None`` retains the native loader. A mapping, including an 

69 empty one, enables distributed balancing with those configuration options. 

70 """ 

71 

72 target: Target[Any] 

73 collate_fn: Optional[Target[Any]] = None 

74 get_batch: Optional[Target[Any]] = None 

75 cost_model: Optional[Target[Any]] = None 

76 balancing_algorithm: Optional[Target[Any]] = None 

77 load_balance: Optional[dict[str, Any]] = None 

78 sampler_type: Literal["single", "cyclic"] = "single" 

79 data_rearrange_map: Any = None 

80 data_sharding: bool = False 

81 use_background_prefetcher: bool = False 

82 

83 def build(self, **runtime_kwargs: Any) -> Any: 

84 """Build the DataLoader target with runtime Dataset arguments.""" 

85 return self.target.build(**runtime_kwargs) 

86 

87 def __getattr__(self, name: str) -> Any: 

88 """Expose configured DataLoader options through the wrapped target.""" 

89 return getattr(self.target, name) 

90 

91 def to_dict(self) -> dict[str, Any]: 

92 """Serialize components in their compact nested YAML shape.""" 

93 config = self.target.to_dict() 

94 config["collate_fn"] = _serialize_config_value(self.collate_fn) 

95 config["get_batch"] = _serialize_config_value(self.get_batch) 

96 config["cost_model"] = _serialize_config_value(self.cost_model) 

97 config["balancing_algorithm"] = _serialize_config_value(self.balancing_algorithm) 

98 if self.load_balance is not None: 

99 config["load_balance"] = _serialize_config_value(self.load_balance) 

100 config["sampler_type"] = self.sampler_type 

101 config["data_rearrange_map"] = _serialize_config_value(self.data_rearrange_map) 

102 config["data_sharding"] = self.data_sharding 

103 config["use_background_prefetcher"] = self.use_background_prefetcher 

104 return config