Coverage for /home/jenkins/.local/lib/python3.10/site-packages/hyper_parallel/data/omni/omni

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

7# http://www.apache.org/licenses/LICENSE-2.0

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"""Omni transform lifecycle and the default AutoProcessor implementation."""

17from __future__ import annotations

19import copy

20import io

21import os

22from collections.abc import Callable, Mapping

23from dataclasses import dataclass

24from typing import Any

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

27from transformers import AutoProcessor

29from hyper_parallel.data.constants import IGNORE_INDEX, OMNI_CP_TOKEN_FIELDS, ONLINE_SOURCE_PATH_KEY

30from hyper_parallel.data.dataset_logging import get_dataset_logger

31from hyper_parallel.data.online.format import (

32 convert_sharegpt_to_messages,

33 normalize_chat_media_fields,

34 resolve_media_path,

35)

37logger = get_dataset_logger(__name__)

40def build_labels_from_supervised_start(

41 input_ids: torch.Tensor,

42 supervised_start: int,

43 *,

44 media_mask: torch.Tensor | None = None,

45) -> torch.Tensor:

46 """Supervise text from a token position through the end of one sample.

48 Args:

49 input_ids: One encoded sample with shape [sequence].

50 supervised_start: First supervised token after the prompt or BOS.

51 media_mask: Media positions excluded from supervision.

53 Returns:

54 Labels retaining supervised text and masking other positions with IGNORE_INDEX.

55 """

56 if input_ids.ndim != 1:

57 raise ValueError("Omni label construction requires one-dimensional input_ids")

59 if not 0 <= supervised_start <= input_ids.numel():

60 raise ValueError(f"Invalid supervised start {supervised_start} for length {input_ids.numel()}")

62 # [prompt..., reply, EOS] -> [-100..., reply, EOS]; next-token shifting belongs to the model loss.

63 labels = input_ids.clone()

64 labels[:supervised_start] = IGNORE_INDEX

66 if media_mask is not None:

67 if media_mask.shape != input_ids.shape:

68 raise ValueError("media_mask must have the same shape as input_ids")

70 labels[media_mask.bool()] = IGNORE_INDEX

72 return labels

75def build_auto_processor(

76 *,

77 pretrained_model_name_or_path: str,

78 trust_remote_code: bool = True,

79 **processor_kwargs: Any,

80) -> Any:

81 """Build the Transformers processor used by the default Omni transform.

83 Args:

84 pretrained_model_name_or_path: Model directory or Hub identifier.

85 trust_remote_code: Whether Transformers may load model-owned code.

86 **processor_kwargs: Additional ``AutoProcessor.from_pretrained`` options.

88 Returns:

89 Processor with an attached tokenizer and resolved chat template.

90 """

91 processor = AutoProcessor.from_pretrained(

92 pretrained_model_name_or_path,

93 trust_remote_code=trust_remote_code,

94 **processor_kwargs,

95 )

96 tokenizer = getattr(processor, "tokenizer", None)

97 if tokenizer is None:

98 raise ValueError("AutoProcessor must provide a tokenizer")

100 chat_template = getattr(processor, "chat_template", None)

101 if chat_template is None:

102 chat_template = getattr(tokenizer, "chat_template", None)

103 if chat_template is not None:

104 processor.chat_template = chat_template

105

106 return processor

107

108

109class OmniDataTransform:

110 """Define model-specific sample and batch encoding hooks.

111

112 Eager transforms implement ``encode_sample`` to return model inputs and labels.

113 Deferred transforms provide ``preencode_sample`` and ``postencode_sample``;

114 offline metadata needs only ``postencode_sample``. The pipeline detects

115 these hook combinations automatically.

116

117 Packing selection and concatenation belong to the optional packing path

118 rather than this transform.

119 """

120

121 def __init__(

122 self, *, max_seq_len: int, processor: Any, training_mode: str = "sft",

123 ) -> None:

124 """Store the sequence limit and externally built processor."""

125 if training_mode not in ("sft", "pretrain"):

126 raise ValueError(f"Unsupported Omni training mode: {training_mode!r}")

127

128 self.training_mode = training_mode

129 if max_seq_len <= 0:

130 raise ValueError("max_seq_len must be positive")

131

132 if processor is None:

133 raise ValueError("processor is required")

134

135 self.max_seq_len = max_seq_len

136 self.processor = processor

137

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

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

140 messages = sample.get("messages")

141 if not isinstance(messages, list):

142 raise ValueError("Omni sample messages must be a list")

143

144 valid_sample = bool(messages)

145 return valid_sample

146

147 def encode_sample(self, sample: dict[str, Any]) -> dict[str, Any] | None:

148 """Encode a complete sample with unshifted labels, or return None to skip it."""

149 del sample

150 raise NotImplementedError(f"{type(self).__name__} does not implement sample encoding")

151

152 def preencode_sample(self, sample: dict[str, Any]) -> dict[str, Any]:

153 """Prepare metadata used by a packing selector without full encoding."""

154 del sample

155 raise NotImplementedError(f"{type(self).__name__} does not implement preencode_sample")

156

157 def postencode_sample(self, sample: dict[str, Any]) -> dict[str, Any]:

158 """Encode selected online or offline metadata into a model-ready sample."""

159 del sample

160 raise NotImplementedError(f"{type(self).__name__} does not implement postencode_sample")

161

162 def encode_batch(self, batch: Any) -> Any:

163 """Apply an optional model-specific batch encoding stage.

164

165 Args:

166 batch: Collated model inputs.

167 """

168 return batch

169

170 def prepare_messages(self, sample: dict[str, Any]) -> list[dict[str, Any]]:

171 """Bind external images in placeholder order and resolve source-relative paths.

172

173 Args:

174 sample: Messages, optional external images, and source path metadata.

175 """

176 messages = sample.get("messages")

177 if not isinstance(messages, list) or not messages:

178 raise ValueError("Omni sample must contain non-empty messages")

179

180 prepared_messages = copy.deepcopy(messages)

181 self._bind_external_images(prepared_messages, sample.get("images"))

182 source_path = sample.get(ONLINE_SOURCE_PATH_KEY)

183 if source_path is None:

184 return prepared_messages

185

186 source_directory = os.path.dirname(os.path.abspath(source_path))

187 for message in prepared_messages:

188 content = message.get("content")

189 if not isinstance(content, list):

190 continue

191 for content_block in content:

192 self._resolve_media_block(content_block, source_directory)

193 return prepared_messages

194

195 def _bind_external_images(self, messages: list, images: Any) -> None:

196 """Consume external images only after checking the complete placeholder count."""

197 if images is None:

198 images = []

199 if not isinstance(images, list):

200 raise TypeError("Omni sample images must be a list")

201 placeholders = []

202 for message in messages:

203 content = message.get("content")

204 if isinstance(content, str) and (images or "<image>" in content):

205 content = [{"type": "text", "text": content}]

206 if isinstance(content, list):

207 message["content"] = self._expand_image_placeholders(content, placeholders)

208 if len(placeholders) != len(images):

209 raise ValueError(f"Image placeholder count {len(placeholders)} does not match images count {len(images)}")

210 for index, (block, image) in enumerate(zip(placeholders, images)):

211 block["image"] = self._decode_external_image(image, index)

212

213 @staticmethod

214 def _expand_image_placeholders(content: list, placeholders: list) -> list:

215 """Preserve text around each placeholder and leave explicit media untouched."""

216 expanded = []

217 for block in content:

218 if not isinstance(block, dict):

219 expanded.append(block)

220 continue

221 text_content = block.get("text")

222 if (block.get("type") == "text"

223 and isinstance(text_content, str) and "<image>" in text_content):

224 for index, text in enumerate(block["text"].split("<image>")):

225 if index:

226 image_block = {"type": "image"}

227 expanded.append(image_block)

228 placeholders.append(image_block)

229 if text:

230 expanded.append({**block, "text": text})

231 else:

232 expanded.append(block)

233 if block.get("type") == "image":

234 # Arrow structs can materialize absent source fields as nulls.

235 source_fields = ("image", "url", "path", "base64")

236 for field in source_fields:

237 if block.get(field) is None:

238 block.pop(field, None)

239 if not any(field in block for field in source_fields):

240 placeholders.append(block)

241 return expanded

242

243 @staticmethod

244 def _decode_external_image(image: Any, index: int) -> Any:

245 """Prefer embedded bytes over path, returning a fully loaded RGB image."""

246 if isinstance(image, str) and image:

247 return image

248 if isinstance(image, Mapping):

249 payload = image.get("bytes")

250 if payload is not None:

251 if not isinstance(payload, (bytes, bytearray, memoryview)):

252 raise TypeError(f"images[{index}].bytes must be binary image data")

253 # Pillow is optional for configurations that only use text/path messages.

254 from PIL import Image # pylint: disable=C0415

255

256 try:

257 with io.BytesIO(payload) as buffer, Image.open(buffer) as decoded:

258 return decoded.convert("RGB")

259 except (OSError, ValueError) as error:

260 raise ValueError(f"Cannot decode images[{index}].bytes") from error

261 path = image.get("path")

262 if isinstance(path, str) and path:

263 return path

264 raise ValueError(f"images[{index}] must be a path or a mapping with bytes/path")

265

266 @classmethod

267 def _resolve_media_block(cls, content_block: Any, source_directory: str) -> None:

268 """Resolve path-bearing fields without changing the processor's message format."""

269 if not isinstance(content_block, dict):

270 return

271

272 block_type = content_block.get("type")

273 if block_type == "image_url":

274 image_url = content_block.get("image_url")

275 if isinstance(image_url, str):

276 content_block["image_url"] = cls._resolve_media_path(image_url, source_directory)

277 elif isinstance(image_url, dict) and isinstance(image_url.get("url"), str):

278 image_url["url"] = cls._resolve_media_path(image_url["url"], source_directory)

279 return

280

281 if block_type in ("image", "video", "audio"):

282 for field in (block_type, "url", "path"):

283 media_field = content_block.get(field)

284 if isinstance(media_field, str):

285 content_block[field] = cls._resolve_media_path(media_field, source_directory)

286

287 @staticmethod

288 def _resolve_media_path(media_path: str, source_directory: str) -> str:

289 """Delegate media path resolution for existing model adapters."""

290 resolved_path = resolve_media_path(media_path, source_directory)

291 return resolved_path

292

293

294# Omni hooks are alternatives: eager encoding intentionally leaves deferred hooks unimplemented.

295class AutoProcessorTransform(OmniDataTransform): # pylint: disable=abstract-method

296 """Encode a conversation and supervise only its final assistant response."""

297

298 def __init__(

299 self, *, max_seq_len: int, processor: Any, training_mode: str = "sft",

300 sample_format: str = "messages",

301 ) -> None:

302 """Select messages/OpenAI/HF or fixed ShareGPT input explicitly."""

303 super().__init__(max_seq_len=max_seq_len, processor=processor, training_mode=training_mode)

304 if sample_format not in ("messages", "openai", "hf", "sharegpt"):

305 raise ValueError(f"Unsupported AutoProcessor sample format: {sample_format!r}")

306

307 self.sample_format = sample_format

308

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

310 """Filter the configured conversation column before source construction."""

311 message_key = "conversations" if self.sample_format == "sharegpt" else "messages"

312 messages = sample.get(message_key)

313 if not isinstance(messages, list):

314 raise ValueError(f"Omni {self.sample_format} sample requires a {message_key} list")

315

316 valid_sample = bool(messages)

317 return valid_sample

318

319 def _assistant_prompt_length(self, messages: list, input_ids: Any, prompt_ids: Any) -> int:

320 """Locate the prompt in processor output without retokenizing media spans.

321

322 A token crossing the source/target character boundary is masked in full.

323 This preserves native processor inputs, unlike separately encoding and

324 concatenating the source and target, which changes sequence-aligned fields.

325 """

326 if len(prompt_ids) < len(input_ids) and input_ids[:len(prompt_ids)].equal(prompt_ids):

327 return len(prompt_ids)

328

329 full_text = self.processor.apply_chat_template(

330 messages, tokenize=False, add_generation_prompt=False,

331 )

332 prompt_text = self.processor.apply_chat_template(

333 messages[:-1], tokenize=False, add_generation_prompt=True,

334 )

335 if not full_text.startswith(prompt_text):

336 raise ValueError("Assistant generation prompt must be a text prefix of the full conversation")

337

338 full_tokens = input_ids.tolist()

339 prompt_tokens = prompt_ids.tolist()

340 common_length = 0

341 for full_token, prompt_token in zip(full_tokens, prompt_tokens):

342 if full_token != prompt_token:

343 break

344 common_length += 1

345 return self._merged_prompt_length(full_tokens, prompt_tokens, common_length)

346

347 def _merged_prompt_length(self, full_tokens: list[int], prompt_tokens: list[int], common_length: int) -> int:

348 """Accept only text-equivalent retokenization after the shared token prefix."""

349 decode_kwargs = {"skip_special_tokens": False, "clean_up_tokenization_spaces": False}

350 tokenizer = self.processor.tokenizer

351 full_tail = tokenizer.decode(full_tokens[common_length:], **decode_kwargs)

352 prompt_tail = tokenizer.decode(prompt_tokens[common_length:], **decode_kwargs)

353 # Only the prompt must decode losslessly to establish the supervision boundary.

354 # A literal U+FFFD in the assistant response is data, not a prompt mismatch.

355 if not prompt_tail or "\ufffd" in prompt_tail or not full_tail.startswith(prompt_tail):

356 raise ValueError("Assistant prompt tokens differ beyond a lossless text boundary merge")

357 for end in range(common_length + 1, len(full_tokens) + 1):

358 decoded = tokenizer.decode(full_tokens[common_length:end], **decode_kwargs)

359 if full_tail.startswith(decoded) and decoded.startswith(prompt_tail):

360 return end

361 raise ValueError("Cannot locate assistant supervision boundary in processor tokens")

362

363 def encode_sample(self, sample: dict[str, Any]) -> dict[str, Any]:

364 """Build model inputs and assistant-only labels from processor messages."""

365 if self.training_mode != "sft":

366 raise NotImplementedError("Image-caption pretraining requires a model-specific encoding method")

367

368 if self.sample_format == "sharegpt":

369 messages = convert_sharegpt_to_messages(sample)

370 else:

371 prepared_sample = dict(sample, messages=self.prepare_messages(sample))

372 messages = normalize_chat_media_fields(prepared_sample)

373 if len(messages) < 2 or messages[-1].get("role") != "assistant":

374 raise ValueError("Omni training sample must end with a final assistant message")

375

376 encoded_sample = dict(self.processor.apply_chat_template(

377 messages,

378 tokenize=True,

379 return_dict=True,

380 return_tensors="pt",

381 add_generation_prompt=False,

382 ))

383 # The generation prompt includes the assistant header: [user, header | reply, EOS].

384 prompt_sample = self.processor.apply_chat_template(

385 messages[:-1],

386 tokenize=True,

387 return_dict=True,

388 return_tensors="pt",

389 add_generation_prompt=True,

390 )

391

392 input_ids = encoded_sample["input_ids"].squeeze(0)

393 prompt_ids = prompt_sample["input_ids"].squeeze(0)

394 if len(input_ids) > self.max_seq_len:

395 raise ValueError("Omni sample exceeds max_seq_len; refusing to truncate a possible media span")

396 prompt_length = self._assistant_prompt_length(messages, input_ids, prompt_ids)

397

398 media_mask = None

399 if "mm_token_type_ids" in encoded_sample:

400 media_mask = encoded_sample["mm_token_type_ids"].squeeze(0) != 0

401 labels = build_labels_from_supervised_start(

402 input_ids, prompt_length, media_mask=media_mask,

403 )

404

405 # Single-sample token fields: [1, L] -> [L]; media tensors keep the processor layout.

406 for field in OMNI_CP_TOKEN_FIELDS.intersection(encoded_sample):

407 encoded_sample[field] = encoded_sample[field].squeeze(0)

408 if not bool(labels[1:].ne(IGNORE_INDEX).any()):

409 raise ValueError("Omni sample has no supervised assistant tokens")

410 encoded_sample["labels"] = labels

411 return encoded_sample

412

413

414SampleEncoder = Callable[[Any], Any]

415

416

417@dataclass(frozen=True)

418class _OmniTransformStrategy:

419 """Resolve the encoding stages implemented by an Omni transform."""

420

421 sample_encoder: SampleEncoder | None

422 selected_sample_encoder: SampleEncoder | None

423

424 @classmethod

425 def from_transform(cls, transform: OmniDataTransform) -> "_OmniTransformStrategy":

426 """Select eager or deferred encoding from the implemented pipeline hooks."""

427 transform_type = type(transform)

428 encode_sample = transform_type.encode_sample is not OmniDataTransform.encode_sample

429 preencode_sample = transform_type.preencode_sample is not OmniDataTransform.preencode_sample

430 postencode_sample = transform_type.postencode_sample is not OmniDataTransform.postencode_sample

431

432 if encode_sample and (preencode_sample or postencode_sample):

433 raise TypeError(

434 f"{type(transform).__name__} cannot combine encode_sample with "

435 "preencode_sample or postencode_sample"

436 )

437

438 if preencode_sample and not postencode_sample:

439 raise TypeError(

440 f"{type(transform).__name__} must pair preencode_sample with postencode_sample"

441 )

442

443 if not encode_sample and not postencode_sample:

444 raise TypeError(

445 f"{type(transform).__name__} must implement encode_sample or postencode_sample"

446 )

447

448 # encode_sample -> encode before candidate selection

449 # preencode_sample + postencode_sample -> prepare metadata, then defer full encoding

450 # postencode_sample only -> preserve offline metadata, then defer encoding

451 sample_encoder = None

452 if preencode_sample:

453 sample_encoder = transform.preencode_sample

454 elif encode_sample:

455 sample_encoder = transform.encode_sample

456

457 selected_sample_encoder = None

458 if postencode_sample:

459 selected_sample_encoder = transform.postencode_sample

460

461 transform_strategy = cls(

462 sample_encoder=sample_encoder,

463 selected_sample_encoder=selected_sample_encoder,

464 )

465 return transform_strategy

466

467 def encode_selected_sample(self, sample: Any) -> Any:

468 """Finish encoding a sample after packing selection.

469

470 Args:

471 sample: Selected candidate.

472 """

473 encoded_sample = sample

474 if self.selected_sample_encoder is not None:

475 encoded_sample = self.selected_sample_encoder(sample)

476 return encoded_sample

477

478 def prepare_samples(self, raw_sample: Any) -> list[Any]:

479 """Return encoded, preencoded, or unchanged samples for candidate selection."""

480 prepared_sample = raw_sample

481 if self.sample_encoder is not None:

482 prepared_sample = self.sample_encoder(raw_sample)

483 if prepared_sample is None:

484 logger.warning("Skipping sample rejected by transform")

485 return []

486

487 return [prepared_sample]

488

489

490__all__ = [

491 "build_auto_processor",

492 "OmniDataTransform",

493 "AutoProcessorTransform",

494]

Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / data / omni / omni_transform.py: 79%

269 statements