Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / data / tools / offline_preparation.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# modified from https://github.com/NVIDIA/Megatron-LM/blob/main/tools/preprocess_data.py 

16 

17"""Processing large data for pretraining.""" 

18 

19import argparse 

20import glob 

21import gzip 

22import json 

23import logging 

24import math 

25import multiprocessing 

26import os 

27import time 

28import traceback 

29import warnings 

30from pathlib import Path 

31from typing import Any 

32 

33import numpy as np 

34from transformers import AutoTokenizer 

35 

36try: 

37 import nltk 

38 from nltk.tokenize.punkt import PunktLanguageVars 

39except ImportError: 

40 PunktLanguageVars = object 

41 NLTK_AVAILABLE = False 

42else: 

43 NLTK_AVAILABLE = True 

44 

45# Store generated samples in the indexed ``.bin/.idx`` format. 

46from hyper_parallel.data.tools import io as indexed_dataset 

47from hyper_parallel.data.dataset_logging import get_dataset_logger 

48 

49logger = get_dataset_logger(__name__) 

50 

51logger = logging.getLogger(__name__) 

52 

53 

54class CustomLanguageVars(PunktLanguageVars): 

55 """Preserve newline runs when Punkt detects sentence boundaries.""" 

56 

57 _period_context_fmt = r""" 

58 \S* 

59 %(SentEndChars)s 

60 \s* 

61 (?=(?P<after_tok> 

62 %(NonWord)s 

63 | 

64 (?P<next_tok>\S+) 

65 ))""" 

66 

67 

68def build_tokenizer(args: argparse.Namespace) -> Any: 

69 """Build the tokenizer configured for offline preprocessing.""" 

70 warnings.warn( 

71 "Only Hugging Face tokenizers loaded through AutoTokenizer are currently supported; custom tokenizer " 

72 "classes are not supported.", 

73 UserWarning, 

74 stacklevel=2, 

75 ) 

76 return build_huggingface_tokenizer(args) 

77 

78 

79def build_huggingface_tokenizer(args: argparse.Namespace) -> Any: 

80 """Build a Hugging Face tokenizer from command-line arguments.""" 

81 tokenizer_path = args.tokenizer_name_or_path 

82 if not tokenizer_path: 

83 raise ValueError("tokenizer_name_or_path must be provided for HuggingFaceTokenizer") 

84 

85 tokenizer_kwargs = { 

86 "use_fast": args.tokenizer_use_fast, 

87 "trust_remote_code": args.trust_remote_code, 

88 } 

89 

90 tokenizer = AutoTokenizer.from_pretrained( 

91 tokenizer_path, 

92 **tokenizer_kwargs, 

93 ) 

94 

95 if args.chat_template: 

96 tokenizer.chat_template = args.chat_template 

97 

98 if args.add_special_tokens: 

99 tokenizer.add_special_tokens( 

100 { 

101 "additional_special_tokens": args.add_special_tokens, 

102 } 

103 ) 

104 

105 return tokenizer 

106 

107 

108def append_eod(args: argparse.Namespace, tokenizer: Any) -> int | None: 

109 """Resolve the document-ending token requested for preprocessing.""" 

110 if not args.append_eod: 

111 return None 

112 eod_id = getattr(tokenizer, "eos_token_id", None) 

113 if eod_id is None: 

114 # Prefer EOS as the document boundary and fall back to SEP for tokenizers 

115 # that do not define an EOS token. 

116 eod_id = getattr(tokenizer, "sep_token_id", None) 

117 if eod_id is None: 

118 raise ValueError("append_eod requires the tokenizer to define eos_token_id or sep_token_id") 

119 return int(eod_id) 

120 

121 

122class Encoder: 

123 """Encode JSON text fields into token ids and sentence lengths.""" 

124 

125 def __init__(self, args: argparse.Namespace) -> None: 

126 """Store the preprocessing arguments shared with worker processes. 

127 

128 Args: 

129 args: Parsed offline preprocessing arguments. 

130 """ 

131 self.args = args 

132 

133 def initializer(self) -> None: 

134 """Build the tokenizer and sentence splitter in each worker process.""" 

135 # Use Encoder class as a container for global data 

136 Encoder.tokenizer = build_tokenizer(self.args) 

137 if self.args.split_sentences: 

138 if not NLTK_AVAILABLE: 

139 raise ImportError("NLTK is required when --split-sentences is enabled") 

140 if os.environ.get("NLTK_DATA"): 

141 library = os.path.join( 

142 os.environ["NLTK_DATA"], 

143 "tokenizers", 

144 "punkt", 

145 f"{self.args.lang}.pickle", 

146 ) 

147 splitter = nltk.load(f"file:{library}") 

148 else: 

149 splitter = nltk.load(f"tokenizers/punkt/{self.args.lang}.pickle") 

150 if self.args.keep_newlines: 

151 Encoder.splitter = nltk.tokenize.punkt.PunktSentenceTokenizer( 

152 train_text=splitter._params, # pylint: disable=protected-access 

153 lang_vars=CustomLanguageVars(), 

154 ) 

155 else: 

156 Encoder.splitter = splitter 

157 

158 def split(self, json_line: str) -> tuple[str, int]: 

159 """Split configured JSON text fields into sentence lists.""" 

160 data = json.loads(json_line) 

161 output = {} 

162 max_chunk_length = 1_000_000 

163 for key in self.args.json_keys: 

164 text = data[key] 

165 sentences_by_chunk = [ 

166 Encoder.splitter.tokenize(text[start : start + max_chunk_length]) 

167 for start in range(0, len(text), max_chunk_length) 

168 ] 

169 output[key] = [sentence for chunk in sentences_by_chunk for sentence in chunk] 

170 return json.dumps(output), len(json_line) 

171 

172 def encode( 

173 self, 

174 json_line: str, 

175 ) -> tuple[dict[str, list[int]], dict[str, list[int]], int]: 

176 """Tokenize one JSON record into per-key token ids and sentence lengths. 

177 

178 Args: 

179 json_line: One serialized JSON record. 

180 

181 Returns: 

182 Per-key token ids, per-key sentence lengths, and the input byte length. 

183 """ 

184 data = json.loads(json_line) 

185 ids = {} 

186 lens = {} 

187 keys = self.args.json_keys 

188 for key in keys: 

189 text = data[key] 

190 if isinstance(text, list): 

191 sentences = text 

192 else: 

193 sentences = [text] 

194 doc_ids = [] 

195 sentence_lens = [] 

196 for sentence in sentences: 

197 sentence_ids = Encoder.tokenizer.encode(sentence, add_special_tokens=False) 

198 if len(sentence_ids) > 0: 

199 doc_ids.extend(sentence_ids) 

200 sentence_lens.append(len(sentence_ids)) 

201 eod = append_eod(self.args, Encoder.tokenizer) 

202 if len(doc_ids) > 0 and eod is not None: 

203 doc_ids.append(eod) 

204 sentence_lens[-1] += 1 

205 ids[key] = doc_ids 

206 lens[key] = sentence_lens 

207 return ids, lens, len(json_line) 

208 

209 

210class Partition: 

211 """Process one dataset partition with a pool of worker processes.""" 

212 

213 def __init__(self, args: argparse.Namespace, workers: int) -> None: 

214 """Store the partition configuration. 

215 

216 Args: 

217 args: Parsed offline preprocessing arguments. 

218 workers: Number of worker processes used by this partition. 

219 """ 

220 self.args = args 

221 self.workers = workers 

222 self.performance: list[float] = [] 

223 

224 def print_processing_stats( 

225 self, 

226 count: int, 

227 proc_start: float, 

228 total_bytes_processed: int, 

229 ) -> None: 

230 """Report encoding throughput at the configured log interval. 

231 

232 Args: 

233 count: Number of documents processed so far. 

234 proc_start: Timestamp when processing started. 

235 total_bytes_processed: Total input bytes processed so far. 

236 """ 

237 if count % self.args.log_interval != 0: 

238 return 

239 elapsed = time.time() - proc_start 

240 docs_per_second = count / elapsed 

241 megabytes_per_second = total_bytes_processed / elapsed / 1024 / 1024 

242 logger.info( 

243 "Processed %d documents (%s docs/s, %s MB/s).", 

244 count, 

245 docs_per_second, 

246 megabytes_per_second, 

247 ) 

248 if self.args.find_optimal_num_workers: 

249 self.performance.append(docs_per_second) 

250 

251 def split_sentences(self, file_name: tuple[str, str]) -> None: 

252 """Split every document in one JSONL partition into sentences.""" 

253 input_file_name, output_file_name = file_name 

254 logger.info("Opening %s", input_file_name) 

255 encoder = Encoder(self.args) 

256 with open(input_file_name, "r", encoding="utf-8") as input_file, open( 

257 output_file_name, 

258 "w", 

259 encoding="utf-8", 

260 ) as output_file: 

261 # multiprocessing.Pool must be shut down with close()+join() so 

262 # in-flight tasks finish; a 'with' block would terminate() them. 

263 pool = multiprocessing.Pool(self.workers, initializer=encoder.initializer) # pylint: disable=R1732 

264 try: 

265 split_docs = pool.imap(encoder.split, input_file, 32) 

266 proc_start = time.time() 

267 total_bytes_processed = 0 

268 for count, (document, bytes_processed) in enumerate(split_docs, start=1): 

269 total_bytes_processed += bytes_processed 

270 output_file.write(document + "\n") 

271 self.print_processing_stats(count, proc_start, total_bytes_processed) 

272 finally: 

273 pool.close() 

274 pool.join() 

275 

276 def _create_builders(self, output_prefix: str, tokenizer: Any) -> tuple[dict[str, str], dict[str, Any]]: 

277 """Create indexed dataset builders for every configured JSON key.""" 

278 level = "sentence" if self.args.split_sentences else "document" 

279 output_idx_files = {} 

280 builders = {} 

281 for key in self.args.json_keys: 

282 output_bin_file = f"{output_prefix}_{key}_{level}.bin" 

283 output_idx_files[key] = f"{output_prefix}_{key}_{level}.idx" 

284 builders[key] = indexed_dataset.IndexedDatasetBuilder( 

285 output_bin_file, 

286 dtype=indexed_dataset.DType.optimal_dtype(len(tokenizer)), 

287 ) 

288 return output_idx_files, builders 

289 

290 def _get_chunk_size(self) -> int | None: 

291 """Return the packed chunk size including the shifted target token.""" 

292 pack_to_seq_len = getattr(self.args, "pack_to_seq_len", None) 

293 return pack_to_seq_len + 1 if pack_to_seq_len is not None else None 

294 

295 def _write_encoded_document( 

296 self, 

297 builders: dict[str, Any], 

298 document: dict[str, list[int]], 

299 sentence_lengths: dict[str, list[int]], 

300 token_buffers: dict[str, list[int]], 

301 chunk_size: int | None, 

302 ) -> None: 

303 """Write one encoded document, optionally packing fixed-size chunks.""" 

304 for key in self.args.json_keys: 

305 if chunk_size is None: 

306 builders[key].add_document(document[key], sentence_lengths[key]) 

307 continue 

308 token_buffers[key].extend(document[key]) 

309 complete_length = len(token_buffers[key]) // chunk_size * chunk_size 

310 for offset in range(0, complete_length, chunk_size): 

311 chunk = token_buffers[key][offset : offset + chunk_size] 

312 builders[key].add_document(chunk, [chunk_size]) 

313 del token_buffers[key][:complete_length] 

314 

315 def _process_encoded_documents( 

316 self, 

317 pool: Any, 

318 encoder: Encoder, 

319 input_file_name: str, 

320 builders: dict[str, Any], 

321 token_buffers: dict[str, list[int]], 

322 chunk_size: int | None, 

323 proc_start: float, 

324 ) -> None: 

325 """Encode input records and write them to the indexed dataset builders.""" 

326 total_bytes_processed = 0 

327 with open(input_file_name, "r", encoding="utf-8") as input_file: 

328 encoded_docs = pool.imap(encoder.encode, input_file, 32) 

329 for count, (document, sentence_lengths, bytes_processed) in enumerate(encoded_docs, start=1): 

330 if self.args.find_optimal_num_workers and count > self.args.max_documents: 

331 break 

332 total_bytes_processed += bytes_processed 

333 self._write_encoded_document( 

334 builders, document, sentence_lengths, token_buffers, chunk_size 

335 ) 

336 self.print_processing_stats(count, proc_start, total_bytes_processed) 

337 

338 def process_json_file(self, file_name: tuple[str, str]) -> list[float]: 

339 """Tokenize one JSONL partition into indexed dataset .bin/.idx files. 

340 

341 Args: 

342 file_name: Input JSONL path and output prefix pair. 

343 

344 Returns: 

345 Throughput measurements collected while benchmarking workers. 

346 """ 

347 input_file_name, output_prefix = file_name 

348 logger.info("Opening %s", input_file_name) 

349 

350 startup_start = time.time() 

351 encoder = Encoder(self.args) 

352 tokenizer = build_tokenizer(self.args) 

353 # multiprocessing.Pool must be shut down with close()+join() so 

354 # in-flight tasks finish; a 'with' block would terminate() them. 

355 pool = multiprocessing.Pool(self.workers, initializer=encoder.initializer) # pylint: disable=R1732 

356 

357 output_idx_files, builders = self._create_builders(output_prefix, tokenizer) 

358 timing = (time.time() - startup_start, time.time()) 

359 logger.info("Time to startup: %s", timing[0]) 

360 chunk_size = self._get_chunk_size() 

361 token_buffers = {key: [] for key in self.args.json_keys} 

362 try: 

363 self._process_encoded_documents( 

364 pool, encoder, input_file_name, builders, token_buffers, chunk_size, timing[1] 

365 ) 

366 for key in self.args.json_keys: 

367 builders[key].finalize(output_idx_files[key]) 

368 finally: 

369 for builder in builders.values(): 

370 builder.data_file.close() 

371 pool.close() 

372 pool.join() 

373 

374 return self.performance 

375 

376 

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

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

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

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

381 return True 

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

383 return False 

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

385 

386 

387def _add_offline_source_arguments(group: Any) -> None: 

388 """Register source, tokenizer, and tokenization arguments.""" 

389 arguments = ( 

390 (("--dataset-name-or-path",), {"required": True, "help": "Input JSON/JSONL file, directory, or glob."}), 

391 (("--output-prefix",), {"required": True, "help": "Path prefix for generated .bin/.idx files."}), 

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

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

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

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

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

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

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

399 (("--keep-newlines",), {"action": "store_true", "help": "Preserve newlines while splitting."}), 

400 (("--lang",), {"default": "english", "help": "Punkt language for sentence splitting."}), 

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

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

403 ) 

404 for flags, options in arguments: 

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

406 

407 

408def _add_offline_worker_arguments(group: Any) -> None: 

409 """Register partitioning and worker arguments.""" 

410 arguments = ( 

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

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

413 (("--workers",), {"type": int, "default": 8, "help": "Number of worker processes."}), 

414 (("--partitions",), {"type": int, "default": 1, "help": "Number of data partitions."}), 

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

416 ( 

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

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

419 ), 

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

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

422 ) 

423 for flags, options in arguments: 

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

425 

426 

427def get_args() -> argparse.Namespace: 

428 """Parse offline dataset preparation arguments.""" 

429 parser = argparse.ArgumentParser(description="Prepare an offline dataset.") 

430 group = parser.add_argument_group(title="input data") 

431 _add_offline_source_arguments(group) 

432 _add_offline_worker_arguments(group) 

433 return parser.parse_args() 

434 

435 

436def get_file_name(args: argparse.Namespace, file_id: int) -> dict[str, str]: 

437 """Build input and output paths for one dataset partition. 

438 

439 Args: 

440 args: Parsed preprocessing arguments. 

441 file_id: Zero-based partition identifier. 

442 

443 Returns: 

444 Input, sentence-split, and output-prefix paths for the partition. 

445 """ 

446 output_directory = os.path.dirname(args.output_prefix) 

447 output_name = os.path.basename(args.output_prefix) 

448 work_directory = os.path.join(output_directory, f".{output_name}_preprocess") 

449 input_directory = os.path.join(work_directory, f"partition{file_id}") 

450 partition_directory = os.path.join(output_directory, f"partition{file_id}") 

451 os.makedirs(input_directory, exist_ok=True) 

452 os.makedirs(partition_directory, exist_ok=True) 

453 input_file_name = os.path.join(input_directory, "input.jsonl") 

454 sentence_split_file = os.path.join(input_directory, "input_ss.jsonl") 

455 output_prefix = os.path.join(partition_directory, output_name) 

456 file_names = {"partition": input_file_name, "sentence_split": sentence_split_file, "output_prefix": output_prefix} 

457 return file_names 

458 

459 

460def check_files_exist( 

461 in_ss_out_names: list[dict[str, str]], 

462 key: str, 

463 num_partitions: int, 

464) -> bool: 

465 """Check whether every partition contains the requested file type.""" 

466 for i in range(num_partitions): 

467 if not os.path.exists(in_ss_out_names[i][key]): 

468 return False 

469 return True 

470 

471 

472def find_optimal_num_workers( 

473 performance: dict[int, list[float]], 

474 partitions: int, 

475) -> None: 

476 """Print candidate worker counts ordered by average throughput.""" 

477 results = [] 

478 for workers, measurements in performance.items(): 

479 if measurements: 

480 results.append((workers, float(np.mean(measurements)))) 

481 if not results: 

482 raise ValueError("No worker performance measurements were collected") 

483 results.sort(key=lambda item: item[1], reverse=True) 

484 logger.info("Worker performance results:") 

485 for position, (workers, average_rate) in enumerate(results, start=1): 

486 logger.info("%d. %d workers: %.4f docs/s", position, workers, average_rate) 

487 best_workers, best_rate = results[0] 

488 logger.info( 

489 "Best configuration: %d total workers (%d per partition), %.4f docs/s.", 

490 best_workers, 

491 best_workers // partitions, 

492 best_rate, 

493 ) 

494 

495 

496def _process_json_file_worker( 

497 args: argparse.Namespace, 

498 workers: int, 

499 name: dict[str, str], 

500 input_key: str, 

501 queue: Any, 

502) -> None: 

503 """Encode one partition and return performance or an error to the parent.""" 

504 try: 

505 partition = Partition(args, workers) 

506 measurements = partition.process_json_file((name[input_key], name["output_prefix"])) 

507 queue.put((True, measurements)) 

508 except Exception: # pylint: disable=W0718 

509 queue.put((False, traceback.format_exc())) 

510 

511 

512def _split_sentences_worker( 

513 args: argparse.Namespace, 

514 workers: int, 

515 name: dict[str, str], 

516 queue: Any, 

517) -> None: 

518 """Split one partition and report completion or an error to the parent.""" 

519 try: 

520 partition = Partition(args, workers) 

521 partition.split_sentences((name["partition"], name["sentence_split"])) 

522 queue.put((True, None)) 

523 except Exception: # pylint: disable=W0718 

524 queue.put((False, traceback.format_exc())) 

525 

526 

527def _wait_for_processes(processes: list[Any], queue: Any) -> list[Any]: 

528 """Collect child results, join processes, and surface child failures.""" 

529 results = [queue.get() for _ in processes] 

530 for process in processes: 

531 process.join() 

532 queue.close() 

533 queue.join_thread() 

534 errors = [payload for succeeded, payload in results if not succeeded] 

535 if errors: 

536 raise RuntimeError("Child preprocessing process failed:\n" + "\n".join(errors)) 

537 return [payload for succeeded, payload in results if succeeded] 

538 

539 

540def _count_jsonl_lines(file_name: str) -> int: 

541 """Count records in a plain or gzip-compressed JSONL file.""" 

542 open_file = gzip.open if file_name.endswith(".gz") else open 

543 open_kwargs = {"mode": "rt", "encoding": "utf-8"} 

544 with open_file(file_name, **open_kwargs) as input_file: 

545 return sum(1 for _ in input_file) 

546 

547 

548def _resolve_input_files(dataset_name_or_path: str) -> list[str]: 

549 """Resolve a JSON input file, directory, or glob into stable file order. 

550 

551 Args: 

552 dataset_name_or_path: Input file, directory, or glob expression. 

553 

554 Returns: 

555 Sorted input JSON/JSONL file paths. 

556 

557 Raises: 

558 ValueError: If the input does not resolve to supported files. 

559 """ 

560 input_path = Path(dataset_name_or_path).expanduser() 

561 supported_suffixes = (".json", ".jsonl", ".json.gz", ".jsonl.gz") 

562 if input_path.is_dir(): 

563 input_files = [ 

564 str(path.resolve()) 

565 for path in input_path.iterdir() 

566 if path.is_file() and path.name.lower().endswith(supported_suffixes) 

567 ] 

568 elif input_path.is_file(): 

569 input_files = [str(input_path.resolve())] 

570 else: 

571 input_files = [ 

572 str(Path(file_name).resolve()) 

573 for file_name in glob.glob(dataset_name_or_path) 

574 if Path(file_name).is_file() and file_name.lower().endswith(supported_suffixes) 

575 ] 

576 input_files.sort() 

577 if not input_files: 

578 raise ValueError( 

579 "dataset_name_or_path must be a JSON/JSONL file, a directory containing " 

580 f"JSON/JSONL files, or a matching glob; received {dataset_name_or_path!r}" 

581 ) 

582 return input_files 

583 

584 

585def _merge_partition_outputs( 

586 args: argparse.Namespace, 

587 in_ss_out_names: list[dict[str, str]], 

588) -> None: 

589 """Merge all encoded partition datasets into final output files.""" 

590 level = "sentence" if args.split_sentences else "document" 

591 tokenizer = build_tokenizer(args) 

592 for key in args.json_keys: 

593 output_bin_file = f"{args.output_prefix}_{key}_{level}.bin" 

594 output_idx_file = f"{args.output_prefix}_{key}_{level}.idx" 

595 builder = indexed_dataset.IndexedDatasetBuilder( 

596 output_bin_file, 

597 dtype=indexed_dataset.DType.optimal_dtype(len(tokenizer)), 

598 ) 

599 for name in in_ss_out_names: 

600 partition_prefix = f"{name['output_prefix']}_{key}_{level}" 

601 builder.add_index(partition_prefix) 

602 builder.finalize(output_idx_file) 

603 

604 

605def _cleanup_intermediate_files( 

606 args: argparse.Namespace, 

607 in_ss_out_names: list[dict[str, str]], 

608 input_files: list[str], 

609) -> None: 

610 """Remove generated JSON partition inputs after successful processing.""" 

611 source_files = {str(Path(file_name).resolve()) for file_name in input_files} 

612 for name in in_ss_out_names: 

613 generated_files = [name["sentence_split"]] 

614 if str(Path(name["partition"]).resolve()) not in source_files: 

615 generated_files.append(name["partition"]) 

616 for file_name in generated_files: 

617 if os.path.isfile(file_name): 

618 os.remove(file_name) 

619 

620 output_path = Path(args.output_prefix).expanduser().resolve() 

621 work_directory = output_path.parent / f".{output_path.name}_preprocess" 

622 if work_directory.is_dir(): 

623 for directory, _, _ in os.walk(work_directory, topdown=False): 

624 try: 

625 os.rmdir(directory) 

626 except OSError: 

627 pass 

628 

629 

630def _validate_preparation_args(args: argparse.Namespace) -> list[int]: 

631 """Validate preprocessing options and return worker candidates.""" 

632 if args.partitions <= 0: 

633 raise ValueError("partitions must be greater than zero") 

634 if args.log_interval <= 0: 

635 raise ValueError("log_interval must be greater than zero") 

636 if args.max_documents <= 0: 

637 raise ValueError("max_documents must be greater than zero") 

638 pack_to_seq_len = getattr(args, "pack_to_seq_len", None) 

639 if pack_to_seq_len is not None and pack_to_seq_len <= 0: 

640 raise ValueError("pack_to_seq_len must be greater than zero") 

641 if pack_to_seq_len is not None and not args.append_eod: 

642 warnings.warn( 

643 "pack_to_seq_len is enabled while append_eod is disabled; packed samples will not contain EOD tokens " 

644 "at original document boundaries.", 

645 UserWarning, 

646 stacklevel=2, 

647 ) 

648 if not args.output_prefix: 

649 raise ValueError("output_prefix must be provided") 

650 if not args.tokenizer_name_or_path: 

651 raise ValueError("tokenizer_name_or_path must be provided") 

652 output_parent = Path(args.output_prefix).expanduser().resolve().parent 

653 output_parent.mkdir(parents=True, exist_ok=True) 

654 worker_candidates = list(args.workers_to_check) if args.find_optimal_num_workers else [args.workers] 

655 if not worker_candidates or any(workers <= 0 for workers in worker_candidates): 

656 raise ValueError("worker counts must be greater than zero") 

657 invalid_workers = [workers for workers in worker_candidates if workers % args.partitions != 0] 

658 if invalid_workers: 

659 raise ValueError(f"Worker counts {invalid_workers} must be divisible by partitions ({args.partitions})") 

660 

661 if args.split_sentences and not NLTK_AVAILABLE: 

662 raise ImportError("NLTK is required when --split-sentences is enabled") 

663 return worker_candidates 

664 

665 

666def _partition_input_data( 

667 args: argparse.Namespace, 

668 input_files: list[str], 

669) -> list[dict[str, str]]: 

670 """Create or reuse input partitions and return their path descriptions.""" 

671 if args.partitions == 1 and len(input_files) == 1: 

672 output_path = Path(args.output_prefix).expanduser().resolve() 

673 work_directory = output_path.parent / f".{output_path.name}_preprocess" 

674 work_directory.mkdir(parents=True, exist_ok=True) 

675 return [ 

676 { 

677 "partition": input_files[0], 

678 "sentence_split": str(work_directory / "input_ss.jsonl"), 

679 "output_prefix": args.output_prefix, 

680 } 

681 ] 

682 

683 partition_size = 1 

684 if args.keep_sequential_samples: 

685 total_sample_count = sum(_count_jsonl_lines(filename) for filename in input_files) 

686 if total_sample_count == 0: 

687 raise ValueError("Input JSON/JSONL files must contain at least one record") 

688 partition_size = math.ceil(total_sample_count / args.partitions) 

689 if args.partitions == 1: 

690 output_path = Path(args.output_prefix).expanduser().resolve() 

691 directory = output_path.parent / f".{output_path.name}_preprocess" / "partition0" 

692 directory.mkdir(parents=True, exist_ok=True) 

693 names = [ 

694 { 

695 "partition": str(directory / "input.jsonl"), 

696 "sentence_split": str(directory / "input_ss.jsonl"), 

697 "output_prefix": args.output_prefix, 

698 } 

699 ] 

700 else: 

701 names = [get_file_name(args, file_id) for file_id in range(args.partitions)] 

702 keep_files = getattr(args, "keep_partition_files", False) 

703 partitions_present = keep_files and check_files_exist(names, "partition", args.partitions) 

704 split_files_present = keep_files and check_files_exist(names, "sentence_split", args.partitions) 

705 if not partitions_present and not split_files_present: 

706 _write_input_partitions(args, names, input_files, partition_size) 

707 return names 

708 

709 

710def _write_input_partitions( 

711 args: argparse.Namespace, 

712 names: list[dict[str, str]], 

713 input_files: list[str], 

714 partition_size: int | None, 

715) -> None: 

716 """Distribute input records across partition files.""" 

717 outputs = [open(name["partition"], "w", encoding="utf-8") for name in names] # pylint: disable=R1732 

718 partition_index = 0 

719 line_count = 0 

720 try: 

721 for input_file_name in input_files: 

722 open_file = gzip.open if input_file_name.endswith(".gz") else open 

723 with open_file(input_file_name, "rt", encoding="utf-8") as input_file: 

724 for line in input_file: 

725 outputs[partition_index].write(line) 

726 if args.keep_sequential_samples: 

727 line_count += 1 

728 if line_count % partition_size == 0 and partition_index < args.partitions - 1: 

729 partition_index += 1 

730 else: 

731 partition_index = (partition_index + 1) % args.partitions 

732 finally: 

733 for output in outputs: 

734 output.close() 

735 

736 

737def _split_partitions(args: argparse.Namespace, workers: int, names: list[dict[str, str]]) -> None: 

738 """Run sentence splitting for partitions that lack cached outputs.""" 

739 split_files_present = getattr(args, "keep_partition_files", False) and check_files_exist( 

740 names, "sentence_split", args.partitions 

741 ) 

742 if not args.split_sentences or split_files_present: 

743 return 

744 queue = multiprocessing.Queue() 

745 processes = [] 

746 for name in names: 

747 process = multiprocessing.Process(target=_split_sentences_worker, args=(args, workers, name, queue)) 

748 process.start()