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

16Inspect token sequences stored in an indexed dataset. 

17 

18python -m hyper_parallel.data.tools.read_indexed_dataset \ 

19 --path outputs/training_demo/tiny_llama_wikitext2/data_no_pack/wikitext2_real_text_document.idx \ 

20 --tokenizer outputs/training_demo/tiny_llama_wikitext2/tokenizer \ 

21 --num-samples 5 

22""" 

23 

24import argparse 

25import logging 

26from collections.abc import Sequence 

27 

28from transformers import AutoTokenizer 

29 

30from hyper_parallel.data.dataset_logging import get_dataset_logger 

31from hyper_parallel.data.indexed.indexed_data_reader import IndexedDataReader 

32 

33logger = get_dataset_logger(__name__) 

34 

35 

36def _path_prefix(path: str) -> str: 

37 """Remove an optional indexed Dataset file suffix from a path.""" 

38 return path[:-4] if path.endswith((".bin", ".idx")) else path 

39 

40 

41def _parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: 

42 """Parse indexed Dataset inspection arguments.""" 

43 parser = argparse.ArgumentParser(description="Inspect an indexed dataset") 

44 parser.add_argument("--path", required=True, help="Dataset prefix or its .bin/.idx path") 

45 parser.add_argument("--tokenizer", help="Local Hugging Face tokenizer path or repository") 

46 parser.add_argument( 

47 "--tokenizer-type", choices=("huggingface", "pretokenized"), default="huggingface", 

48 help="Use pretokenized when only vocabulary metadata is available", 

49 ) 

50 parser.add_argument("--vocab-size", type=int, help="Vocabulary size for pretokenized data") 

51 parser.add_argument("--eod-token-id", type=int, help="EOD token ID for pretokenized data") 

52 parser.add_argument("--start-index", type=int, default=0, help="First sequence index to read") 

53 parser.add_argument("--num-samples", type=int, default=5, help="Number of sequences to read") 

54 parser.add_argument("--max-token-ids", type=int, default=80, help="Maximum token IDs printed per sequence") 

55 parser.add_argument("--max-text-chars", type=int, default=1000, help="Maximum decoded characters per sequence") 

56 parser.add_argument("--skip-special-tokens", action="store_true", help="Remove special tokens when decoding") 

57 return parser.parse_args(argv) 

58 

59 

60def main(argv: Sequence[str] | None = None) -> None: 

61 """Print indexed Dataset metadata and selected token sequences.""" 

62 args = _parse_args(argv) 

63 if args.start_index < 0 or args.num_samples < 0: 

64 raise ValueError("start-index and num-samples must be non-negative") 

65 

66 if args.max_token_ids < 0 or args.max_text_chars < 0: 

67 raise ValueError("max-token-ids and max-text-chars must be non-negative") 

68 

69 reader = IndexedDataReader(_path_prefix(args.path)) 

70 tokenizer = None 

71 if args.tokenizer_type == "huggingface" and args.tokenizer: 

72 tokenizer = AutoTokenizer.from_pretrained(args.tokenizer) 

73 

74 if args.tokenizer_type == "pretokenized": 

75 if args.vocab_size is None or args.eod_token_id is None: 

76 raise ValueError("pretokenized mode requires vocab-size and eod-token-id") 

77 logger.info( 

78 "Pretokenized metadata: vocab_size=%d, eod_token_id=%d; decoded text is unavailable", 

79 args.vocab_size, args.eod_token_id, 

80 ) 

81 sequence_lengths = reader.sequence_lengths 

82 logger.info( 

83 "Dataset: sequences=%d, documents=%d, min_length=%d, max_length=%d, mean_length=%.2f", 

84 len(reader), len(reader.document_indices) - 1, int(sequence_lengths.min()), int(sequence_lengths.max()), 

85 float(sequence_lengths.mean()), 

86 ) 

87 

88 stop_index = min(args.start_index + args.num_samples, len(reader)) 

89 for index in range(args.start_index, stop_index): 

90 token_ids = reader[index].tolist() 

91 displayed_ids = token_ids[:args.max_token_ids] 

92 suffix = " ..." if len(displayed_ids) < len(token_ids) else "" 

93 logger.info("[%d] length=%d\ntoken_ids=%s%s", index, len(token_ids), displayed_ids, suffix) 

94 

95 if args.tokenizer_type == "pretokenized": 

96 eod_positions = [position for position, token_id in enumerate(token_ids) if token_id == args.eod_token_id] 

97 logger.info("eod_positions=%s", eod_positions) 

98 

99 if tokenizer is not None: 

100 text = tokenizer.decode(token_ids, skip_special_tokens=args.skip_special_tokens) 

101 text_suffix = " ..." if len(text) > args.max_text_chars else "" 

102 logger.info("text=%r%s", text[:args.max_text_chars], text_suffix) 

103 

104 

105if __name__ == "__main__": 

106 logging.basicConfig(level=logging.INFO, format="%(message)s") 

107 main()