Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / data / tools / offline_config.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 

16"""Configuration for Hugging Face offline dataset preparation.""" 

17 

18import argparse 

19import re 

20from dataclasses import dataclass, field 

21from pathlib import Path 

22from typing import Dict, List, Optional, Union 

23 

24 

25_DatasetDataFiles = Union[ 

26 str, 

27 List[str], 

28 Dict[str, Union[str, List[str]]], 

29] 

30 

31 

32@dataclass 

33class OfflinePreparationConfig: 

34 """Values required to download and preprocess one Hugging Face dataset.""" 

35 

36 dataset_name_or_path: str 

37 output_prefix: str 

38 tokenizer_name_or_path: str 

39 download_dir: Optional[str] = None 

40 dataset_subset_name: Optional[str] = None 

41 dataset_split: str = "train" 

42 revision: Optional[str] = None 

43 cache_dir: Optional[str] = None 

44 data_dir: Optional[str] = None 

45 data_files: Optional[_DatasetDataFiles] = None 

46 num_proc: Optional[int] = None 

47 json_keys: Union[str, List[str]] = field(default_factory=lambda: ["text"]) 

48 tokenizer_use_fast: bool = True 

49 trust_remote_code: bool = False 

50 chat_template: Optional[str] = None 

51 add_special_tokens: Optional[List[str]] = None 

52 split_sentences: bool = False 

53 keep_newlines: bool = False 

54 lang: str = "english" 

55 workers: int = 8 

56 partitions: int = 1 

57 append_eod: bool = True 

58 pack_to_seq_len: Optional[int] = None 

59 keep_sequential_samples: bool = True 

60 keep_partition_files: bool = False 

61 find_optimal_num_workers: bool = False 

62 workers_to_check: List[int] = field(default_factory=lambda: [16, 32, 64]) 

63 max_documents: int = 100_000 

64 log_interval: int = 1000 

65 

66 @staticmethod 

67 def _slug(value: str) -> str: 

68 """Normalize a repository, subset, or tokenizer name for file paths.""" 

69 normalized = re.sub(r"[^A-Za-z0-9._-]+", "-", value).strip("-._") 

70 if not normalized: 

71 raise ValueError(f"Could not derive a file name from {value!r}") 

72 return normalized 

73 

74 def dataset_label(self) -> str: 

75 """Derive a local dataset label from the subset or repository ID. 

76 

77 Returns: 

78 File-system-safe dataset label. 

79 """ 

80 source = ( 

81 self.dataset_subset_name 

82 or self.dataset_name_or_path.rsplit("/", maxsplit=1)[-1] 

83 ) 

84 return self._slug(source) 

85 

86 def json_keys_list(self) -> List[str]: 

87 """Return validated JSON key names as a list. 

88 

89 Returns: 

90 JSON keys consumed by the offline encoder. 

91 

92 Raises: 

93 ValueError: If no valid JSON key is configured. 

94 """ 

95 keys = ( 

96 [self.json_keys] 

97 if isinstance(self.json_keys, str) 

98 else list(self.json_keys) 

99 ) 

100 if not keys or any( 

101 not isinstance(key, str) or not key.strip() 

102 for key in keys 

103 ): 

104 raise ValueError("json_keys must contain one or more column names") 

105 return keys 

106 

107 def download_root(self) -> Path: 

108 """Resolve the root directory for downloaded raw JSONL datasets. 

109 

110 Returns: 

111 Custom download directory, or the default 

112 ``./download_datasets/{dataset_label}/``. 

113 """ 

114 if self.download_dir and self.download_dir.strip(): 

115 return Path(self.download_dir).expanduser().resolve() 

116 return ( 

117 Path("download_datasets") / self.dataset_label() 

118 ).expanduser().resolve() 

119 

120 def resolved_json_path(self) -> Path: 

121 """Resolve the configured or derived local JSONL path. 

122 

123 Returns: 

124 Local JSONL destination path. 

125 """ 

126 source_path = Path(self.dataset_name_or_path).expanduser() 

127 local_suffixes = (".json", ".jsonl", ".json.gz", ".jsonl.gz") 

128 if source_path.is_dir(): 

129 return source_path.resolve() 

130 if source_path.name.lower().endswith(local_suffixes): 

131 resolved_path = source_path.resolve() 

132 if not resolved_path.is_file(): 

133 raise FileNotFoundError( 

134 f"Local JSON dataset does not exist: {resolved_path}" 

135 ) 

136 return resolved_path 

137 

138 file_name = f"{self.dataset_label()}-{self._slug(self.dataset_split)}.jsonl" 

139 return self.download_root() / file_name 

140 

141 def resolved_output_prefix(self) -> Path: 

142 """Resolve the indexed dataset output prefix. 

143 

144 Returns: 

145 Output prefix without the ``.bin/.idx`` suffixes. 

146 """ 

147 if not self.output_prefix.strip(): 

148 raise ValueError("output_prefix must be a non-empty path") 

149 return Path(self.output_prefix).expanduser().resolve() 

150 

151 def to_offline_args(self) -> argparse.Namespace: 

152 """Convert this typed config to offline preprocessing arguments. 

153 

154 Returns: 

155 Complete preprocessing arguments for ``prepare_offline_dataset``. 

156 """ 

157 return argparse.Namespace( 

158 dataset_name_or_path=str(self.resolved_json_path()), 

159 output_prefix=str(self.resolved_output_prefix()), 

160 json_keys=self.json_keys_list(), 

161 tokenizer_name_or_path=self.tokenizer_name_or_path, 

162 tokenizer_use_fast=self.tokenizer_use_fast, 

163 trust_remote_code=self.trust_remote_code, 

164 chat_template=self.chat_template, 

165 add_special_tokens=self.add_special_tokens, 

166 split_sentences=self.split_sentences, 

167 keep_newlines=self.keep_newlines, 

168 lang=self.lang, 

169 append_eod=self.append_eod, 

170 pack_to_seq_len=self.pack_to_seq_len, 

171 keep_sequential_samples=self.keep_sequential_samples, 

172 keep_partition_files=self.keep_partition_files, 

173 workers=self.workers, 

174 partitions=self.partitions, 

175 find_optimal_num_workers=self.find_optimal_num_workers, 

176 workers_to_check=self.workers_to_check, 

177 max_documents=self.max_documents, 

178 log_interval=self.log_interval, 

179 )