Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / data / tools / huggingface_offline.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"""Download a Hugging Face dataset and convert it to indexed binary data. 

17 

18This module is both importable and runnable as a CLI entry point:: 

19 

20 python -m hyper_parallel.data.tools.huggingface_offline \\ 

21 --dataset Salesforce/wikitext \\ 

22 --dataset-subset wikitext-103-raw-v1 \\ 

23 --output-prefix ./offline_datasets/my_dataset \\ 

24 --tokenizer gpt2 \\ 

25 --json-keys text \\ 

26 --workers 8 

27""" 

28 

29import argparse 

30import os 

31from pathlib import Path 

32from typing import Any, List 

33 

34from hyper_parallel.data.dataset_logging import get_dataset_logger 

35from hyper_parallel.data.tools.offline_config import OfflinePreparationConfig 

36from hyper_parallel.data.tools.offline_preparation import ( 

37 prepare_offline_dataset, 

38) 

39 

40logger = get_dataset_logger(__name__) 

41 

42 

43def _download_jsonl(config: OfflinePreparationConfig) -> Path: 

44 """Download and normalize the configured Hugging Face split as JSONL.""" 

45 json_path = config.resolved_json_path() 

46 source_path = Path(config.dataset_name_or_path).expanduser() 

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

48 if source_path.is_dir() or source_path.name.lower().endswith(local_suffixes): 

49 logger.info("Using local JSON dataset input at %s", json_path) 

50 return json_path 

51 

52 if json_path.is_file(): 

53 logger.info("Reusing downloaded Hugging Face dataset at %s", json_path) 

54 return json_path 

55 

56 logger.info( 

57 "Downloading Hugging Face dataset %s (config=%s, split=%s)", 

58 config.dataset_name_or_path, 

59 config.dataset_subset_name, 

60 config.dataset_split, 

61 ) 

62 from datasets import load_dataset # pylint: disable=C0415 

63 

64 if config.num_proc is not None and config.num_proc <= 0: 

65 raise ValueError("num_proc must be greater than zero") 

66 

67 load_dataset_kwargs: dict[str, Any] = { 

68 "path": config.dataset_name_or_path, 

69 "split": config.dataset_split, 

70 } 

71 optional_load_dataset_kwargs = { 

72 "name": config.dataset_subset_name, 

73 "revision": config.revision, 

74 "cache_dir": config.cache_dir, 

75 "data_dir": config.data_dir, 

76 "data_files": config.data_files, 

77 "num_proc": config.num_proc, 

78 } 

79 load_dataset_kwargs.update({key: value for key, value in optional_load_dataset_kwargs.items() if value is not None}) 

80 dataset = load_dataset(**load_dataset_kwargs) 

81 keys = config.json_keys_list() 

82 missing_keys = [key for key in keys if key not in dataset.column_names] 

83 if missing_keys: 

84 raise ValueError( 

85 f"Dataset {config.dataset_name_or_path} does not contain configured " 

86 f"keys {missing_keys}; " 

87 f"available columns: {dataset.column_names}" 

88 ) 

89 dataset = dataset.select_columns(keys) 

90 

91 json_path.parent.mkdir(parents=True, exist_ok=True) 

92 dataset.to_json( 

93 str(json_path), 

94 orient="records", 

95 lines=True, 

96 force_ascii=False, 

97 ) 

98 logger.info("Saved Hugging Face dataset to %s", json_path) 

99 return json_path 

100 

101 

102def _parse_bool(value: str) -> bool: 

103 """Parse an explicit command-line boolean value.""" 

104 normalized = value.strip().lower() 

105 if normalized in {"true", "1", "yes", "on"}: 

106 return True 

107 if normalized in {"false", "0", "no", "off"}: 

108 return False 

109 raise argparse.ArgumentTypeError(f"Expected a boolean value, but received {value!r}") 

110 

111 

112def _add_huggingface_arguments(parser: argparse.ArgumentParser) -> None: 

113 """Register Hugging Face source and output arguments.""" 

114 dataset_group = parser.add_argument_group("dataset") 

115 dataset_arguments = ( 

116 (("--dataset",), {"required": True, "help": "Dataset ID or local JSON/JSONL path."}), 

117 (("--dataset-subset",), {"default": None, "help": "Optional dataset subset."}), 

118 (("--dataset-split",), {"default": "train", "help": "Dataset split."}), 

119 (("--revision",), {"default": None, "help": "Optional dataset revision."}), 

120 (("--cache-dir",), {"default": None, "help": "Optional dataset cache directory."}), 

121 (("--data-dir",), {"default": None, "help": "Optional repository data directory."}), 

122 (("--data-files",), {"nargs": "+", "default": None, "help": "Optional source data files."}), 

123 (("--num-proc",), {"type": int, "default": None, "help": "Dataset preparation process count."}), 

124 (("--json-keys",), {"nargs": "+", "default": ["text"], "help": "JSON fields to tokenize."}), 

125 ) 

126 for flags, options in dataset_arguments: 

127 dataset_group.add_argument(*flags, **options) 

128 output_group = parser.add_argument_group("output") 

129 output_group.add_argument("--output-prefix", required=True, help="Generated .bin/.idx path prefix.") 

130 output_group.add_argument("--download-dir", default=None, help="Raw JSONL download directory.") 

131 

132 

133def _add_conversion_arguments(parser: argparse.ArgumentParser) -> None: 

134 """Register tokenizer, preprocessing, parallelism, and benchmark arguments.""" 

135 groups = { 

136 "tokenizer": ( 

137 (("--tokenizer",), {"required": True, "help": "Tokenizer name or local path."}), 

138 (("--tokenizer-use-fast",), {"type": _parse_bool, "default": True, "help": "Use fast tokenizer."}), 

139 (("--trust-remote-code",), {"action": "store_true", "help": "Allow tokenizer remote code."}), 

140 (("--chat-template",), {"default": None, "help": "Optional tokenizer chat template."}), 

141 (("--add-special-tokens",), {"nargs": "+", "default": None, "help": "Additional special tokens."}), 

142 ), 

143 "preprocessing": ( 

144 (("--split-sentences",), {"action": "store_true", "help": "Split text into sentences."}), 

145 (("--keep-newlines",), {"action": "store_true", "help": "Preserve newline runs."}), 

146 (("--lang",), {"default": "english", "help": "Punkt language."}), 

147 (("--append-eod",), {"type": _parse_bool, "default": True, "help": "Append an EOD token."}), 

148 (("--pack-to-seq-len",), {"type": int, "default": None, "help": "Fixed packed sequence length."}), 

149 ), 

150 "parallelism": ( 

151 (("--workers",), {"type": int, "default": 8, "help": "Worker process count."}), 

152 (("--partitions",), {"type": int, "default": 1, "help": "Data partition count."}), 

153 (("--keep-sequential-samples",), {"action": "store_true", "help": "Keep sample order."}), 

154 (("--keep-partition-files",), {"action": "store_true", "help": "Keep partition files."}), 

155 ), 

156 "benchmark": ( 

157 (("--find-optimal-num-workers",), {"action": "store_true", "help": "Benchmark worker counts."}), 

158 ( 

159 ("--workers-to-check",), 

160 {"nargs": "+", "type": int, "default": [16, 32, 64], "help": "Candidate worker counts."}, 

161 ), 

162 (("--max-documents",), {"type": int, "default": 100_000, "help": "Benchmark document limit."}), 

163 (("--log-interval",), {"type": int, "default": 1000, "help": "Progress-report interval."}), 

164 ), 

165 } 

166 for title, arguments in groups.items(): 

167 group = parser.add_argument_group(title) 

168 for flags, options in arguments: 

169 group.add_argument(*flags, **options) 

170 

171 

172def _get_args(argv: List[str] | None = None) -> argparse.Namespace: 

173 """Parse the Hugging Face offline preparation CLI arguments.""" 

174 parser = argparse.ArgumentParser(description="Download a Hugging Face dataset and convert to .bin/.idx") 

175 _add_huggingface_arguments(parser) 

176 _add_conversion_arguments(parser) 

177 return parser.parse_args(argv) 

178 

179 

180def main(argv: List[str] | None = None) -> None: 

181 """Download, tokenize, and convert a dataset to ``.bin/.idx`` files.""" 

182 args = _get_args(argv) 

183 

184 config = OfflinePreparationConfig( 

185 dataset_name_or_path=args.dataset, 

186 dataset_subset_name=args.dataset_subset, 

187 dataset_split=args.dataset_split, 

188 revision=args.revision, 

189 cache_dir=args.cache_dir, 

190 data_dir=args.data_dir, 

191 data_files=args.data_files, 

192 num_proc=args.num_proc, 

193 json_keys=args.json_keys, 

194 output_prefix=args.output_prefix, 

195 download_dir=args.download_dir, 

196 tokenizer_name_or_path=args.tokenizer, 

197 tokenizer_use_fast=args.tokenizer_use_fast, 

198 trust_remote_code=args.trust_remote_code, 

199 chat_template=args.chat_template, 

200 add_special_tokens=args.add_special_tokens, 

201 split_sentences=args.split_sentences, 

202 keep_newlines=args.keep_newlines, 

203 lang=args.lang, 

204 workers=args.workers, 

205 partitions=args.partitions, 

206 append_eod=args.append_eod, 

207 pack_to_seq_len=args.pack_to_seq_len, 

208 keep_sequential_samples=args.keep_sequential_samples, 

209 keep_partition_files=args.keep_partition_files, 

210 find_optimal_num_workers=args.find_optimal_num_workers, 

211 workers_to_check=args.workers_to_check, 

212 max_documents=args.max_documents, 

213 log_interval=args.log_interval, 

214 ) 

215 

216 if int(os.environ.get("RANK", "0")) != 0: 

217 return 

218 

219 _ = _download_jsonl(config) 

220 prepare_offline_dataset(config.to_offline_args()) 

221 

222 

223if __name__ == "__main__": 

224 main()