Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed_data / indexed_text.py: 82%

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"""Built-in constructor callbacks for unpacked indexed text samples.""" 

16# This distributed-data package is intentionally PyTorch-only. 

17 

18from __future__ import annotations 

19 

20from collections.abc import Mapping, Sequence 

21from typing import Any 

22 

23import torch # pylint: disable=forbidden-backend-import 

24 

25_IGNORE_INDEX = -100 

26 

27 

28def _as_1d_long_tensor(value: Any, field: str) -> torch.Tensor: 

29 """Normalize one source field to a CPU int64 vector.""" 

30 try: 

31 tensor = torch.as_tensor(value, dtype=torch.int64) 

32 except (TypeError, ValueError, RuntimeError) as exc: 

33 raise ValueError(f"Indexed source field {field!r} cannot be converted to an int64 tensor") from exc 

34 if tensor.ndim != 1 or tensor.numel() < 1: 

35 raise ValueError(f"Indexed source field {field!r} must be a non-empty one-dimensional sequence") 

36 return tensor 

37 

38 

39def _validate_pad_token_id(pad_token_id: int) -> None: 

40 if isinstance(pad_token_id, bool) or not isinstance(pad_token_id, int) or pad_token_id < 0: 

41 raise ValueError("pad_token_id must be a non-negative integer") 

42 

43 

44def split_packed_text_sample( 

45 sample: Mapping[str, Any], *, pad_token_id: int | None = None, 

46) -> tuple[dict[str, torch.Tensor], ...]: 

47 """Split a packed record at explicit independent-attention boundaries. 

48 

49 Args: 

50 sample: Pre-shifted ``input_ids`` and ``labels`` plus leading-zero 

51 ``cu_seq_lens`` covering the complete record. Labels and token order 

52 within each real sequence are preserved, including ignored targets. 

53 pad_token_id: Dedicated source PAD ID. When supplied, remove only the 

54 record's trailing run of this ID with ignored labels. Leave source 

55 records unchanged when omitted. 

56 

57 Returns: 

58 Ordered sequence views, ready for per-sequence metadata estimation. 

59 """ 

60 input_ids = _as_1d_long_tensor(sample["input_ids"], "input_ids") 

61 labels = _as_1d_long_tensor(sample["labels"], "labels") 

62 boundaries = torch.as_tensor(sample["cu_seq_lens"]) 

63 if input_ids.shape != labels.shape: 

64 raise ValueError("Packed input_ids and labels must have identical shapes") 

65 if boundaries.ndim != 1 or boundaries.dtype not in (torch.int32, torch.int64) or boundaries.numel() < 2: 

66 raise ValueError("Packed cu_seq_lens must be an integer boundary vector") 

67 offsets = boundaries.tolist() 

68 if offsets[0] != 0 or offsets[-1] != input_ids.numel() or any( 

69 start >= end for start, end in zip(offsets[:-1], offsets[1:]) 

70 ): 

71 raise ValueError("Packed cu_seq_lens must strictly increase from zero to the record length") 

72 valid_length = input_ids.numel() 

73 if pad_token_id is not None: 

74 _validate_pad_token_id(pad_token_id) 

75 # Ignored prompt targets are real context; labels alone cannot identify PAD. 

76 non_padding = ((input_ids != pad_token_id) | (labels != _IGNORE_INDEX)).nonzero().flatten() 

77 if not non_padding.numel(): 

78 raise ValueError("Packed record must contain at least one non-padding token") 

79 valid_length = int(non_padding[-1]) + 1 

80 return tuple( 

81 {"input_ids": input_ids[start:min(end, valid_length)], "labels": labels[start:min(end, valid_length)]} 

82 for start, end in zip(offsets[:-1], offsets[1:]) 

83 if start < valid_length 

84 ) 

85 

86 

87def pack_indexed_text_samples( 

88 samples: Sequence[Any], seq_len: int, *, pad_token_id: int = 0, 

89) -> dict[str, torch.Tensor]: 

90 """Pack one planned source-sample bin into a fixed-length text row. 

91 

92 Args: 

93 samples: Source mappings containing pre-shifted ``input_ids`` and 

94 ``labels`` vectors. 

95 seq_len: Physical output length for one packed row. 

96 pad_token_id: Token ID used for the independent padding sequence. 

97 

98 Returns: 

99 One padded row and its exact source-sample boundaries. 

100 """ 

101 if not samples: 

102 raise ValueError("Indexed text packing requires at least one source sample") 

103 if not isinstance(seq_len, int) or isinstance(seq_len, bool) or seq_len < 1: 

104 raise ValueError(f"seq_len must be a positive integer, but got {seq_len!r}") 

105 _validate_pad_token_id(pad_token_id) 

106 

107 input_parts = [] 

108 label_parts = [] 

109 sample_lengths = [] 

110 for sample in samples: 

111 if not isinstance(sample, Mapping): 

112 raise ValueError(f"Indexed text samples must be mappings, but got {type(sample)}") 

113 try: 

114 input_ids = _as_1d_long_tensor(sample["input_ids"], "input_ids") 

115 labels = _as_1d_long_tensor(sample["labels"], "labels") 

116 except KeyError as exc: 

117 raise ValueError(f"Indexed text sample is missing required field {exc.args[0]!r}") from exc 

118 if input_ids.shape != labels.shape: 

119 raise ValueError("Indexed text input_ids and labels must have identical shapes") 

120 input_parts.append(input_ids) 

121 label_parts.append(labels) 

122 sample_lengths.append(input_ids.numel()) 

123 

124 packed_length = sum(sample_lengths) 

125 if packed_length > seq_len: 

126 raise ValueError(f"Indexed text bin has {packed_length} tokens, exceeding seq_len={seq_len}") 

127 input_ids = torch.cat(input_parts) 

128 labels = torch.cat(label_parts) 

129 padding_length = seq_len - packed_length 

130 if padding_length: 

131 input_ids = torch.cat((input_ids, input_ids.new_full((padding_length,), pad_token_id))) 

132 labels = torch.cat((labels, labels.new_full((padding_length,), _IGNORE_INDEX))) 

133 

134 sequence_ends = torch.tensor(sample_lengths, dtype=torch.int32).cumsum(dim=0, dtype=torch.int32) 

135 if padding_length: 

136 sequence_ends = torch.cat((sequence_ends, sequence_ends.new_tensor([seq_len]))) 

137 cu_seq_lens = torch.cat((sequence_ends.new_zeros(1), sequence_ends)) 

138 return { 

139 "input_ids": input_ids, 

140 "labels": labels, 

141 "cu_seq_lens": cu_seq_lens, 

142 } 

143 

144 

145def collate_indexed_text_sequences(packed_sequences: Sequence[Any]) -> dict[str, torch.Tensor]: 

146 """Combine fixed-length packed rows into one rank-local training batch. 

147 

148 Args: 

149 packed_sequences: Outputs from :func:`pack_indexed_text_samples`. 

150 

151 Returns: 

152 Batched token fields and flattened global sequence boundaries. 

153 """ 

154 if not packed_sequences: 

155 raise ValueError("Indexed text collation requires at least one packed sequence") 

156 if any(not isinstance(sequence, Mapping) for sequence in packed_sequences): 

157 raise ValueError("Every indexed packed sequence must be a mapping") 

158 

159 input_ids = torch.stack([sequence["input_ids"] for sequence in packed_sequences]) 

160 labels = torch.stack([sequence["labels"] for sequence in packed_sequences]) 

161 if input_ids.shape != labels.shape or input_ids.ndim != 2: 

162 raise ValueError("Indexed packed input_ids and labels must form equal two-dimensional batches") 

163 

164 seq_len = input_ids.shape[1] 

165 boundary_parts = [input_ids.new_zeros(1, dtype=torch.int32)] 

166 for row_index, sequence in enumerate(packed_sequences): 

167 cu_seq_lens = sequence["cu_seq_lens"] 

168 if not isinstance(cu_seq_lens, torch.Tensor) or cu_seq_lens.dtype != torch.int32: 

169 raise ValueError("Indexed packed cu_seq_lens must be an int32 tensor") 

170 if cu_seq_lens.ndim != 1 or cu_seq_lens.numel() < 2 or int(cu_seq_lens[0]) != 0: 

171 raise ValueError("Indexed packed cu_seq_lens must be a leading-zero boundary vector") 

172 if int(cu_seq_lens[-1]) != seq_len: 

173 raise ValueError("Indexed packed cu_seq_lens must cover the complete padded row") 

174 boundary_parts.append(cu_seq_lens[1:] + row_index * seq_len) 

175 return { 

176 "input_ids": input_ids, 

177 "labels": labels, 

178 "cu_seq_lens": torch.cat(boundary_parts), 

179 } 

180 

181 

182__all__ = ["collate_indexed_text_sequences", "pack_indexed_text_samples", "split_packed_text_sample"]