Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / data / text / text_transform.py: 79%
131 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07:18 +0800
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."""
17from __future__ import annotations
19import json
20from collections.abc import Callable, Mapping, Sequence
21from dataclasses import dataclass
22from typing import Any, Literal
24import torch
26from hyper_parallel.data.constants import IGNORE_INDEX
27from hyper_parallel.data.dataset_logging import get_dataset_logger
28from hyper_parallel.data.online.format import InstructionMessages, normalize_conversation
30TextDataType = Literal["plaintext", "sharegpt", "alpaca"]
31logger = get_dataset_logger(__name__)
34def _get_record_value(sample: Mapping[str, Any], keys: str | Sequence[str]) -> Any:
35 if isinstance(keys, str):
36 try:
37 return sample[keys]
38 except KeyError as exc:
39 raise ValueError(f"Sample does not contain field {keys!r}") from exc
40 for key in keys:
41 if key in sample:
42 return sample[key]
43 raise ValueError(f"Sample does not contain any configured text fields: {list(keys)!r}")
46def _normalize_tools(tools: Any, tools_key: str) -> list[dict[str, Any]] | None:
47 """Normalize an explicitly configured tools field for native chat templates."""
48 if tools is None or tools is False:
49 return None
50 if isinstance(tools, str):
51 if not tools:
52 return None
53 try:
54 tools = json.loads(tools)
55 except json.JSONDecodeError as exc:
56 raise ValueError(f"Tools field {tools_key!r} must contain valid JSON") from exc
57 if tools is None or tools is False:
58 return None
59 if not isinstance(tools, Sequence) or isinstance(tools, (str, bytes)):
60 raise ValueError(f"Tools field {tools_key!r} must be a sequence of mappings")
62 normalized_tools = []
63 for tool in tools:
64 if not isinstance(tool, Mapping):
65 raise ValueError(f"Every entry in tools field {tools_key!r} must be a mapping")
66 normalized_tools.append(dict(tool))
67 return normalized_tools or None
70def _encode_messages(
71 chat_template: Any,
72 messages: Sequence[Mapping[str, Any]],
73 max_seq_len: int,
74 tools: Sequence[Mapping[str, Any]] | None = None,
75) -> dict[str, Any]:
76 """Encode normalized messages into a shifted causal-language-model sample."""
77 encode_kwargs = {"max_seq_len": max_seq_len}
78 if tools:
79 encode_kwargs["tools"] = tools
80 encoded = chat_template.encode_messages(messages, **encode_kwargs)
81 input_ids = torch.tensor(encoded["input_ids"], dtype=torch.int64)
82 labels = torch.tensor(encoded["labels"], dtype=torch.int64)
83 return {
84 "input_ids": input_ids[:-1],
85 "labels": labels[1:],
86 }
89class IdentityDataTransform:
90 """Return each input sample unchanged."""
92 def __init__(self, tokenizer: Any = None, chat_template: Any = None) -> None:
93 """Retain optional upstream assets for target compatibility.
95 Args:
96 tokenizer: Optional tokenizer built by the LLM Trainer.
97 chat_template: Optional chat template built from model assets.
98 """
99 self.tokenizer = tokenizer
100 self.chat_template = chat_template
102 @staticmethod
103 def __call__(sample: Any) -> Any:
104 """Return the input sample without modification."""
105 return sample
108@dataclass
109class PlaintextTransform:
110 """Tokenize plaintext records into one or more model samples."""
112 tokenizer: Any
113 max_seq_len: int
114 text_keys: str | Sequence[str] = "text"
115 drop_ragged_tail: bool = False
117 def __post_init__(self) -> None:
118 """Validate the tokenizer and sequence length configuration."""
119 if self.tokenizer is None:
120 raise ValueError("tokenizer is required for plaintext data")
121 if self.max_seq_len <= 0:
122 raise ValueError("max_seq_len must be positive")
124 def __call__(self, sample: Mapping[str, Any]) -> list[dict[str, Any]]:
125 """Tokenize and chunk one plaintext record."""
126 text = _get_record_value(sample, self.text_keys)
127 token_ids = self.tokenizer.encode(text, add_special_tokens=False)
128 eos_token_id = getattr(self.tokenizer, "eos_token_id", None)
129 if eos_token_id is not None:
130 token_ids = [*token_ids, eos_token_id]
132 transformed = []
133 for start in range(0, len(token_ids) - 1, self.max_seq_len):
134 window = token_ids[start:start + self.max_seq_len + 1]
135 if (self.drop_ragged_tail and start > 0
136 and len(window) < self.max_seq_len + 1):
137 # Opt-in for models whose packed-sample boundaries must stay
138 # aligned, such as the V4.1 CSA compression ratio and the
139 # context-parallel shard size: a short tail window violates
140 # both once it lands in a batch. Off by default, so every
141 # other model keeps the trailing partial window.
142 break
143 text = torch.tensor(window, dtype=torch.long)
144 model_sample = {
145 "input_ids": text[:-1],
146 "labels": text[1:],
147 }
148 transformed.append(model_sample)
149 return transformed
151 def is_valid_sample(self, sample: Mapping[str, Any]) -> bool:
152 """Return whether one source record contains non-empty plaintext."""
153 text = _get_record_value(sample, self.text_keys)
154 if not isinstance(text, str):
155 raise ValueError("Plaintext sample text must be a string")
156 is_valid = text.strip() != ""
157 return is_valid
160@dataclass
161class TextConversationTransform:
162 """Encode conversation records with a configured chat template."""
164 chat_template: Any
165 max_seq_len: int
166 text_keys: str | Sequence[str] = "conversation"
167 tools_key: str | None = None
168 message_formatter: Callable[[Mapping[str, Any]], list[dict[str, Any]]] | None = None
170 def __post_init__(self) -> None:
171 """Validate the chat template and sequence length configuration."""
172 if self.chat_template is None:
173 raise ValueError("chat_template is required for conversation data")
174 if self.max_seq_len <= 0:
175 raise ValueError("max_seq_len must be positive")
176 if self.tools_key is not None and (not isinstance(self.tools_key, str) or not self.tools_key):
177 raise ValueError("tools_key must be a non-empty string or None")
179 def __call__(self, sample: Mapping[str, Any]) -> list[dict[str, Any]]:
180 """Encode one conversation record."""
181 messages = (self.message_formatter(sample) if self.message_formatter is not None
182 else normalize_conversation(_get_record_value(sample, self.text_keys)))
183 tools = None
184 if self.tools_key is not None:
185 tools = _normalize_tools(sample.get(self.tools_key), self.tools_key)
186 encoded = _encode_messages(self.chat_template, messages, self.max_seq_len, tools)
187 return [encoded] if bool(encoded["labels"].ne(IGNORE_INDEX).any()) else []
189 def is_valid_sample(self, sample: Mapping[str, Any]) -> bool:
190 """Return whether one source record contains conversation messages."""
191 messages = (self.message_formatter(sample) if self.message_formatter is not None
192 else _get_record_value(sample, self.text_keys))
193 is_valid = bool(messages)
194 return is_valid
197def build_text_transform(
198 data_type: TextDataType,
199 *,
200 tokenizer: Any = None,
201 chat_template: Any = None,
202 max_seq_len: int,
203 text_keys: str | Sequence[str] | None = None,
204 drop_ragged_tail: bool = False,
205 tools_key: str | None = None,
206 column_mapping: Mapping[str, str] | None = None,
207 input_separator: str = "\n\n",
208) -> Callable[[Any], Any]:
209 """Build the transform selected by the text dataset format.
211 Args:
212 data_type: One of plaintext, sharegpt, or alpaca.
213 tokenizer: Tokenizer used by plaintext transforms.
214 chat_template: Chat template used by conversation transforms.
215 max_seq_len: Maximum model sequence length.
216 text_keys: Source field(s); defaults to text or conversations by format.
217 drop_ragged_tail: Drop trailing partial plaintext windows.
218 tools_key: Optional field containing tool definitions for conversation data.
219 column_mapping: Logical instruction/input/output fields mapped to source columns.
220 input_separator: Separator inserted between non-empty instruction and input fields.
222 Returns:
223 The configured text sample transform.
225 Raises:
226 ValueError: If ``data_type`` is unsupported.
227 """
228 if text_keys is None:
229 text_keys = {"sharegpt": "conversations"}.get(data_type, "text")
230 if data_type == "plaintext":
231 if tools_key is not None:
232 raise ValueError("tools_key is only supported for sharegpt data")
233 data_transform = PlaintextTransform(tokenizer, max_seq_len, text_keys, drop_ragged_tail)
234 elif data_type == "sharegpt":
235 data_transform = TextConversationTransform(chat_template, max_seq_len, text_keys, tools_key)
236 elif data_type == "alpaca":
237 if tools_key is not None:
238 raise ValueError("tools_key is only supported for sharegpt data")
239 if column_mapping is None:
240 raise ValueError("column_mapping is required for alpaca data")
241 data_transform = TextConversationTransform(
242 chat_template,
243 max_seq_len,
244 message_formatter=InstructionMessages(column_mapping, input_separator),
245 )
246 else:
247 raise ValueError(f"Unsupported text data format: {data_type!r}")
249 logger.debug("Built text transform: data_type=%s, transform=%s", data_type, type(data_transform).__name__)
250 return data_transform