Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / data / text / text_transform.py: 60%

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"""Transform plaintext and conversation records into model samples.""" 

16 

17from __future__ import annotations 

18 

19from collections.abc import Callable, Mapping, Sequence 

20from dataclasses import dataclass 

21from typing import Any, Literal 

22 

23import torch 

24 

25from hyper_parallel.data.constants import IGNORE_INDEX 

26from hyper_parallel.data.dataset_logging import get_dataset_logger 

27 

28TextDataType = Literal["plaintext", "conversation"] 

29logger = get_dataset_logger(__name__) 

30 

31 

32def _get_record_value(sample: Mapping[str, Any], keys: str | Sequence[str]) -> Any: 

33 if isinstance(keys, str): 

34 try: 

35 return sample[keys] 

36 except KeyError as exc: 

37 raise ValueError(f"Sample does not contain field {keys!r}") from exc 

38 for key in keys: 

39 if key in sample: 

40 return sample[key] 

41 raise ValueError(f"Sample does not contain any configured text fields: {list(keys)!r}") 

42 

43 

44class IdentityDataTransform: 

45 """Return each input sample unchanged.""" 

46 

47 def __init__(self, tokenizer: Any = None, chat_template: Any = None) -> None: 

48 """Retain optional upstream assets for target compatibility. 

49 

50 Args: 

51 tokenizer: Optional tokenizer built by the LLM Trainer. 

52 chat_template: Optional chat template built from model assets. 

53 """ 

54 self.tokenizer = tokenizer 

55 self.chat_template = chat_template 

56 

57 @staticmethod 

58 def __call__(sample: Any) -> Any: 

59 """Return the input sample without modification.""" 

60 return sample 

61 

62 

63@dataclass 

64class PlaintextTransform: 

65 """Tokenize plaintext records into one or more model samples.""" 

66 

67 tokenizer: Any 

68 max_seq_len: int 

69 text_keys: str | Sequence[str] = "text" 

70 drop_ragged_tail: bool = False 

71 

72 def __post_init__(self) -> None: 

73 """Validate the tokenizer and sequence length configuration.""" 

74 if self.tokenizer is None: 

75 raise ValueError("tokenizer is required for plaintext data") 

76 if self.max_seq_len <= 0: 

77 raise ValueError("max_seq_len must be positive") 

78 

79 def __call__(self, sample: Mapping[str, Any]) -> list[dict[str, Any]]: 

80 """Tokenize and chunk one plaintext record.""" 

81 text = _get_record_value(sample, self.text_keys) 

82 token_ids = self.tokenizer.encode(text, add_special_tokens=False) 

83 eos_token_id = getattr(self.tokenizer, "eos_token_id", None) 

84 if eos_token_id is not None: 

85 token_ids = [*token_ids, eos_token_id] 

86 

87 transformed = [] 

88 for start in range(0, len(token_ids) - 1, self.max_seq_len): 

89 window = token_ids[start:start + self.max_seq_len + 1] 

90 if (self.drop_ragged_tail and start > 0 

91 and len(window) < self.max_seq_len + 1): 

92 # Opt-in for models whose packed-sample boundaries must stay 

93 # aligned, such as the V4.1 CSA compression ratio and the 

94 # context-parallel shard size: a short tail window violates 

95 # both once it lands in a batch. Off by default, so every 

96 # other model keeps the trailing partial window. 

97 break 

98 text = torch.tensor(window, dtype=torch.long) 

99 model_sample = { 

100 "input_ids": text[:-1], 

101 "labels": text[1:], 

102 } 

103 transformed.append(model_sample) 

104 return transformed 

105 

106 def is_valid_sample(self, sample: Mapping[str, Any]) -> bool: 

107 """Return whether one source record contains non-empty plaintext.""" 

108 text = _get_record_value(sample, self.text_keys) 

109 if not isinstance(text, str): 

110 raise ValueError("Plaintext sample text must be a string") 

111 is_valid = text.strip() != "" 

112 return is_valid 

113 

114 

115@dataclass 

116class TextConversationTransform: 

117 """Encode conversation records with a configured chat template.""" 

118 

119 chat_template: Any 

120 max_seq_len: int 

121 text_keys: str | Sequence[str] = "conversation" 

122 

123 def __post_init__(self) -> None: 

124 """Validate the chat template and sequence length configuration.""" 

125 if self.chat_template is None: 

126 raise ValueError("chat_template is required for conversation data") 

127 if self.max_seq_len <= 0: 

128 raise ValueError("max_seq_len must be positive") 

129 

130 def __call__(self, sample: Mapping[str, Any]) -> list[dict[str, Any]]: 

131 """Encode one conversation record.""" 

132 messages = _get_record_value(sample, self.text_keys) 

133 encoded = self.chat_template.encode_messages(messages, max_seq_len=self.max_seq_len) 

134 input_ids = torch.as_tensor(encoded["input_ids"], dtype=torch.long) 

135 labels = torch.as_tensor(encoded["labels"], dtype=torch.long) 

136 shifted_labels = labels[1:] 

137 if not bool(shifted_labels.ne(IGNORE_INDEX).any()): 

138 return [] 

139 

140 model_sample = { 

141 "input_ids": input_ids[:-1], 

142 "labels": shifted_labels, 

143 } 

144 return [model_sample] 

145 

146 def is_valid_sample(self, sample: Mapping[str, Any]) -> bool: 

147 """Return whether one source record contains conversation messages.""" 

148 messages = _get_record_value(sample, self.text_keys) 

149 is_valid = bool(messages) 

150 return is_valid 

151 

152 

153def build_text_transform( 

154 data_type: TextDataType, 

155 *, 

156 tokenizer: Any = None, 

157 chat_template: Any = None, 

158 max_seq_len: int, 

159 text_keys: str | Sequence[str] = "text", 

160 drop_ragged_tail: bool = False, 

161) -> Callable[[Any], Any]: 

162 """Build the transform selected by the text data type. 

163 

164 Args: 

165 data_type: Plaintext or conversation input format. 

166 tokenizer: Tokenizer used by plaintext transforms. 

167 chat_template: Chat template used by conversation transforms. 

168 max_seq_len: Maximum model sequence length. 

169 text_keys: Field or candidate fields containing the source text. 

170 drop_ragged_tail: Drop the trailing partial window of a multi-window 

171 plaintext record instead of emitting it as a shorter sample. 

172 

173 Returns: 

174 The configured text sample transform. 

175 

176 Raises: 

177 ValueError: If ``data_type`` is unsupported. 

178 """ 

179 if data_type == "plaintext": 

180 data_transform = PlaintextTransform( 

181 tokenizer, max_seq_len, text_keys, drop_ragged_tail) 

182 elif data_type == "conversation": 

183 data_transform = TextConversationTransform(chat_template, max_seq_len, text_keys) 

184 else: 

185 raise ValueError(f"Unsupported text data type: {data_type!r}") 

186 

187 logger.debug("Built text transform: data_type=%s, transform=%s", data_type, type(data_transform).__name__) 

188 return data_transform