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

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 tokenizer: Optional[Target[Any]] = None 

33 

34 

35@dataclass 

36class DatasetConfig: 

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

38 

39 target: Target[Any] 

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

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

42 

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

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

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

46 return dataset 

47 

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

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

50 return getattr(self.target, name) 

51 

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

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

54 config = self.target.to_dict() 

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

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

57 return config 

58 

59 

60@dataclass 

61class DataLoaderConfig: 

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

63 

64 target: Target[Any] 

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

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

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

68 data_rearrange_map: Any = None 

69 data_sharding: bool = False 

70 use_background_prefetcher: bool = False 

71 

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

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

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

75 

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

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

78 return getattr(self.target, name) 

79 

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

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

82 config = self.target.to_dict() 

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

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

85 config["sampler_type"] = self.sampler_type 

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

87 config["data_sharding"] = self.data_sharding 

88 config["use_background_prefetcher"] = self.use_background_prefetcher 

89 return config