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

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"""Chat-template implementations used by LLM conversation transforms.

17Canonical merge (05 §11.3) of the former

18``components/data/chat_template.py`` and

19``components/datasets/llm/chat_template.py`` duplicates: one

20``CHAT_TEMPLATE_REGISTRY``, the ``DatasetLogger`` logging path, and

21``IGNORE_INDEX`` shared from the constants module.

22"""

24from abc import ABC, abstractmethod

25from typing import TYPE_CHECKING, Any, Callable, Dict, Iterator, List, MutableMapping, Optional, Sequence, Type, Union

27from hyper_parallel.data.constants import IGNORE_INDEX

28from hyper_parallel.data.dataset_logging import get_dataset_logger

30if TYPE_CHECKING:

31 from transformers import PreTrainedTokenizer

33logger = get_dataset_logger(__name__)

36class Registry(MutableMapping):

37 """Domain registry for chat-template formatters.

39 Inlined from the former ``auto_models/components/utils/registry.py`` in

40 stage 7 (05 §10.3 row 565): the generic Registry was merged into the

41 domain registries that consume it instead of living in a utils package.

42 """

44 # Class instance object, so that a call to `register` can be reflected into all other files correctly, even if

45 # a new instance is created (in order to locally override a given function)

46 registry = []

48 def __init__(self, name: str) -> None:

49 """Initialize the registry.

51 Args:

52 name: Human-readable name of the registered category, used in

53 error messages.

54 """

55 self._name = name

56 self.registry.append(name)

57 self._local_mapping = {}

58 self._global_mapping = {}

60 def __getitem__(self, key: str) -> Union[Type, Callable]:

61 """Look up a registered class or function by key.

63 Args:

64 key: Registration key.

66 Returns:

67 The locally overridden value when present, otherwise the globally

68 registered value.

70 Raises:

71 ValueError: If ``key`` is not registered.

72 """

73 # First check if instance has a local override

74 if key not in self.valid_keys():

75 raise ValueError(f"Unknown {self._name} name: {key}. No {self._name} registered for this source.")

76 if key in self._local_mapping:

77 return self._local_mapping[key]

78 return self._global_mapping[key]

80 def __setitem__(self, key: str, value: Union[Type, Callable]) -> None:

81 """Set a local override for ``key`` without affecting other instances."""

82 # Allow local update of the default functions without impacting other instances

83 self._local_mapping.update({key: value})

85 def __delitem__(self, key: str) -> None:

86 """Delete the local override for ``key``."""

87 del self._local_mapping[key]

89 def __iter__(self) -> Iterator[str]:

90 """Iterate over all valid keys, local overrides taking precedence."""

91 # Ensure we use all keys, with the overwritten ones on top

92 return iter({**self._global_mapping, **self._local_mapping})

94 def __len__(self) -> int:

95 """Return the number of distinct registered keys."""

96 return len(self._global_mapping.keys() | self._local_mapping.keys())

98 def register(

99 self,

100 key: str,

101 cls_or_func: Optional[Union[Type, Callable]] = None,

102 ) -> Union[Type, Callable]:

103 """Register a class or function under ``key``.

104

105 Can be used either as a plain call or as a decorator::

106

107 registry.register("name", MyClass)

108

109 @registry.register("name")

110 class MyClass: ...

111

112 Args:

113 key: Registration key.

114 cls_or_func: The class or function to register. When ``None``, a

115 decorator is returned instead.

116

117 Returns:

118 The registered class/function, or a decorator when ``cls_or_func``

119 is ``None``.

120

121 Raises:

122 ValueError: If ``key`` is already registered (decorator form).

123 """

124 if cls_or_func is not None:

125 self._global_mapping[key] = cls_or_func

126 return cls_or_func

127

128 def decorator(cls_or_func: Union[Type, Callable]) -> Union[Type, Callable]:

129 """Register the decorated class or function under ``key``."""

130 if key in self._global_mapping:

131 raise ValueError(

132 f"{self._name} for '{key}' is already registered. Cannot register duplicate {self._name}."

133 )

134 self._global_mapping.update({key: cls_or_func})

135 return cls_or_func

136

137 return decorator

138

139 def valid_keys(self) -> List[str]:

140 """Return the list of all registered keys."""

141 return list(self.keys())

142

143

144CHAT_TEMPLATE_REGISTRY = Registry("ChatTemplate")

145

146

147def _format_janus_message(

148 message: Dict[str, str], index: int, message_count: int, task_type: str, assistant_count: int

149) -> tuple[str, int]:

150 """Format one Janus message and update the generation assistant count."""

151 role = message["role"]

152 content = message["content"]

153 separators = ["\n\n", "<｜end▁of▁sentence｜>"]

154 if content == "":

155 return role + ":", assistant_count

156 if "assistant" in role and (

157 "wikihow_generation" in task_type or "interleave_generation" in task_type

158 ):

159 prefix = "Assistant: " if assistant_count == 0 else ""

160 suffix = separators[1] if index + 1 == message_count else separators[0]

161 return prefix + content.strip() + suffix, assistant_count + 1

162 if "assistant" in role:

163 return "Assistant: " + content.strip() + separators[1], assistant_count

164 if "user" in role:

165 return "User: " + content.strip() + separators[0], assistant_count

166 if "system" in role and "wikihow_generation" in task_type:

167 instruction = "Please generate a step-by-step tutorial with images for the following question."

168 return content.strip() + separators[0] + instruction + separators[0], assistant_count

169 if "system" in role:

170 return content.strip() + separators[0], assistant_count

171 raise ValueError(f"Unknown role {role}, should be one of {{system, user, assistant}}.")

172

173

174def _janus_labels(content_ids: List[int], image_token_id: int, loss_mask: int, task_type: str) -> List[int]:

175 """Build labels for one encoded Janus message."""

176 if loss_mask != 1:

177 return [IGNORE_INDEX] * len(content_ids)

178 if (

179 image_token_id in content_ids

180 and "wikihow_generation" not in task_type

181 and "interleave_generation" not in task_type

182 ):

183 return [image_token_id if token == image_token_id else IGNORE_INDEX for token in content_ids]

184 return content_ids

185

186

187def build_chat_template(template_name: str, tokenizer: "PreTrainedTokenizer") -> "ChatTemplate":

188 """Build the registered chat template for a tokenizer.

189

190 Args:

191 template_name: Registry key of the chat template implementation.

192 tokenizer: Hugging Face tokenizer used for encoding and special tokens.

193

194 Returns:

195 The configured chat template instance.

196 """

197 return CHAT_TEMPLATE_REGISTRY[template_name](tokenizer)

198

199

200class ChatTemplate(ABC):

201 """

202 Abstract class for chat template.

203 """

204

205 def __init__(self, tokenizer: Any) -> None:

206 """Store the tokenizer used for encoding and special tokens."""

207 self.tokenizer = tokenizer

208

209 def save_pretrained(self, output_dir: str) -> None:

210 """Attach the jinja template to the tokenizer and save it.

211

212 Args:

213 output_dir: Directory the tokenizer assets are written to.

214 """

215 self.tokenizer.chat_template = self.get_jinja_template()

216 try:

217 self.tokenizer.save_pretrained(output_dir)

218 except OSError:

219 logger.warning("Failed to save tokenizer.")

220

221 @abstractmethod

222 def encode_messages(self, messages: Sequence[Dict[str, str]], max_seq_len: int = 8192) -> Dict[str, List[int]]:

223 """

224 Encodes messages to a dictionary of input_ids, attention_mask, and labels.

225 """

226

227 @abstractmethod

228 def get_jinja_template(self) -> str:

229 """

230 Gets the jinja template for the chat template.

231 """

232

233

234@CHAT_TEMPLATE_REGISTRY.register("default")

235class DefaultTemplate(ChatTemplate):

236 """Render each message as ``Role: content`` followed by the EOS token."""

237

238 def encode_messages(self, messages: Sequence[Dict[str, str]], max_seq_len: int = 8192) -> Dict[str, List[int]]:

239 """Encode messages into input_ids, attention_mask, and loss-masked labels.

240

241 Args:

242 messages: Conversation records with ``role``, ``content``, and ``loss_mask`` fields.

243 max_seq_len: Maximum sequence length kept from the end of the encoding.

244

245 Returns:

246 Model inputs with input_ids, attention_mask, and labels.

247 """

248 input_ids, attention_mask, labels = [], [], []

249 for message in messages:

250 content_str = message["role"].title() + ": " + message["content"].strip() + self.tokenizer.eos_token + "\n"

251 content_ids = self.tokenizer.encode(content_str, add_special_tokens=False)

252 input_ids += content_ids

253 attention_mask += [1] * len(content_ids)

254 if message["loss_mask"] == 1:

255 labels += content_ids

256 else:

257 labels += [IGNORE_INDEX] * len(content_ids)

258

259 model_inputs = {"input_ids": input_ids, "attention_mask": attention_mask, "labels": labels}

260 model_inputs = {k: v[-max_seq_len:] for k, v in model_inputs.items()}

261 return model_inputs

262

263 def get_jinja_template(self) -> str:

264 """Return the jinja template matching the default ``Role: content`` rendering."""

265 return (

266 "{% if not add_generation_prompt is defined %}{% set add_generation_prompt = false %}{% endif %}"

267 "{% for message in messages %}"

268 "{{ message['role'].title() + ': ' + message['content'] | trim + eos_token + '\n' }}"

269 "{% endfor %}"

270 "{% if add_generation_prompt %}{{ 'Assistant: ' }}{% endif %}"

271 )

272

273

274@CHAT_TEMPLATE_REGISTRY.register("tokenizer")

275class TokenizerTemplate(ChatTemplate):

276 """Use the tokenizer's native chat template and train only selected messages."""

277

278 def _update_prefix_labels(

279 self,

280 previous_ids: List[int],

281 current_ids: List[int],

282 labels: List[int],

283 ) -> None:

284 """Ensure appending a message does not rewrite the existing token prefix."""

285 # ``labels`` is only adjusted by overrides that permit a terminal rewrite.

286 del labels

287 previous_length = len(previous_ids)

288 if current_ids[:previous_length] != previous_ids:

289 raise ValueError(

290 "The tokenizer chat template structurally rewrote an earlier conversation prefix; "

291 "the generic tokenizer template requires prefix-stable rendering."

292 )

293

294 def encode_messages(

295 self,

296 messages: Sequence[Dict[str, str]],

297 max_seq_len: int = 8192,

298 ) -> Dict[str, List[int]]:

299 """Encode messages with the tokenizer chat template and mask non-assistant loss.

300

301 Each message is appended incrementally so token boundaries follow the

302 tokenizer's own prefix-stable rendering.

303

304 Args:

305 messages: Conversation records with ``role``, ``content``, and optional ``loss_mask``.

306 max_seq_len: Maximum sequence length kept from the end of the encoding.

307

308 Returns:

309 Model inputs with input_ids, attention_mask, and labels.

310

311 Raises:

312 ValueError: If the template rewrites or shortens an earlier conversation prefix.

313 """

314 input_ids: List[int] = []

315 labels: List[int] = []

316 previous_length = 0

317

318 for end, message in enumerate(messages, start=1):

319 encoded = self.tokenizer.apply_chat_template(

320 messages[:end],

321 tokenize=True,

322 add_generation_prompt=False,

323 return_dict=True,

324 )

325 current_ids = encoded["input_ids"]

326 current_length = len(current_ids)

327 if current_length < previous_length:

328 raise ValueError(

329 "The tokenizer chat template shortened the conversation after adding a message; "

330 "assistant-only loss masking requires monotonic message boundaries."

331 )

332

333 self._update_prefix_labels(input_ids, current_ids, labels)

334 loss_mask = message.get("loss_mask", 1 if message["role"] == "assistant" else 0)

335 new_ids = current_ids[previous_length:]

336 labels.extend(new_ids if loss_mask == 1 else [IGNORE_INDEX] * len(new_ids))

337 input_ids = current_ids

338 previous_length = current_length

339

340 input_ids = input_ids[-max_seq_len:]

341 labels = labels[-max_seq_len:]

342 return {

343 "input_ids": input_ids,

344 "attention_mask": [1] * len(input_ids),

345 "labels": labels,

346 }

347

348 def get_jinja_template(self) -> str:

349 """Return the tokenizer's native jinja chat template.

350

351 Raises:

352 ValueError: If the tokenizer does not define a chat template.

353 """

354 if not self.tokenizer.chat_template:

355 raise ValueError("The tokenizer does not define a native chat template.")

356 return self.tokenizer.chat_template

357

358

359@CHAT_TEMPLATE_REGISTRY.register("gpt_oss")

360class GptOssTokenizerTemplate(TokenizerTemplate):

361 """Support GPT-OSS terminal rewrites from <|return|> to <|end|>."""

362

363 def __init__(self, tokenizer: "PreTrainedTokenizer") -> None:

364 """Resolve the GPT-OSS terminal token IDs required for prefix rewrites.

365

366 Args:

367 tokenizer: Hugging Face tokenizer that must define <|return|> and <|end|>.

368

369 Raises:

370 ValueError: If either terminal token is missing from the vocabulary.

371 """

372 super().__init__(tokenizer)

373 self.return_token_id = tokenizer.convert_tokens_to_ids("<|return|>")

374 self.end_token_id = tokenizer.convert_tokens_to_ids("<|end|>")

375 if tokenizer.unk_token_id in (self.return_token_id, self.end_token_id):

376 raise ValueError("The GPT-OSS chat template requires <|return|> and <|end|> tokenizer tokens.")

377

378 def _update_prefix_labels(

379 self,

380 previous_ids: List[int],

381 current_ids: List[int],

382 labels: List[int],

383 ) -> None:

384 """Accept only the terminal <|return|>-to-<|end|> rewrite and fix its label.

385

386 Args:

387 previous_ids: Token IDs rendered before the latest message was appended.

388 current_ids: Token IDs rendered after the latest message was appended.

389 labels: Loss labels aligned with ``previous_ids``, updated in place.

390

391 Raises:

392 ValueError: If the rewrite is not the supported terminal substitution.

393 """

394 previous_length = len(previous_ids)

395 rewritten_positions = [

396 index

397 for index in range(previous_length)

398 if previous_ids[index] != current_ids[index]

399 ]

400 if not rewritten_positions:

401 return

402

403 is_terminal_rewrite = (

404 rewritten_positions == [previous_length - 1]

405 and len(current_ids) > previous_length

406 and previous_ids[-1] == self.return_token_id

407 and current_ids[previous_length - 1] == self.end_token_id

408 and self.return_token_id not in current_ids[previous_length:]

409 )

410 if not is_terminal_rewrite:

411 raise ValueError(

412 "The GPT-OSS tokenizer chat template structurally rewrote an earlier conversation prefix; "

413 "only the terminal <|return|>-to-<|end|> substitution is supported."

414 )

415

416 if labels[-1] != IGNORE_INDEX:

417 labels[-1] = self.end_token_id

418

419

420@CHAT_TEMPLATE_REGISTRY.register("llama2")

421class Llama2Template(ChatTemplate):

422 """Render messages with the Llama-2 [INST]/<<SYS>> prompt format."""

423

424 def encode_messages(self, messages: Sequence[Dict[str, str]], max_seq_len: int = 8192) -> Dict[str, List[int]]:

425 """Encode messages into input_ids, attention_mask, and loss-masked labels.

426

427 Args:

428 messages: Conversation records with ``role``, ``content``, and ``loss_mask`` fields.

429 max_seq_len: Maximum sequence length kept from the end of the encoding.

430

431 Returns:

432 Model inputs with input_ids, attention_mask, and labels.

433

434 Raises:

435 ValueError: If a message role is not one of system, user, assistant, or tool.

436 """

437 input_ids, attention_mask, labels = [], [], []

438 for message in messages:

439 if message["role"] == "system":

440 content_str = "<<SYS>>\n" + message["content"].strip() + "\n<</SYS>>\n\n"

441 elif message["role"] == "user":

442 content_str = self.tokenizer.bos_token + "[INST] " + message["content"].strip() + " [/INST]"

443 elif message["role"] == "assistant":

444 content_str = " " + message["content"].strip() + " " + self.tokenizer.eos_token

445 elif message["role"] == "tool":

446 content_str = self.tokenizer.bos_token + "[TOOL] " + message["content"].strip() + " [/TOOL]"

447 else:

448 raise ValueError(

449 f"Unknown role {message['role']}, should be one of {{system, user, assistant, tool}}."

450 )

451

452 content_ids = self.tokenizer.encode(content_str, add_special_tokens=False)

453 input_ids += content_ids

454 attention_mask += [1] * len(content_ids)

455 if message["loss_mask"] == 1:

456 labels += content_ids

457 else:

458 labels += [IGNORE_INDEX] * len(content_ids)

459

460 model_inputs = {"input_ids": input_ids, "attention_mask": attention_mask, "labels": labels}

461 model_inputs = {k: v[-max_seq_len:] for k, v in model_inputs.items()}

462 return model_inputs

463

464 def get_jinja_template(self) -> str:

465 """Return the jinja template matching the Llama-2 prompt format."""

466 return (

467 "{% if messages[0]['role'] == 'system' %}"

468 "{{ '<<SYS>>\n' + messages[0]['content'] | trim + '\n<</SYS>>\n\n' }}"

469 "{% set loop_messages = messages[1:] %}"

470 "{% else %}"

471 "{% set loop_messages = messages %}"

472 "{% endif %}"

473 "{% for message in loop_messages %}"

474 "{% set content = message['content'] %}"

475 "{% if message['role'] == 'user' %}"

476 "{{ bos_token + '[INST] ' + content | trim + ' [/INST]' }}"

477 "{% elif message['role'] == 'tool' %}"

478 "{{ bos_token + '[TOOL] ' + content | trim + ' [/TOOL]' }}"

479 "{% elif message['role'] == 'assistant' %}"

480 "{{ ' ' + content | trim + ' ' + eos_token }}"

481 "{% endif %}"

482 "{% endfor %}"

483 )

484

485

486@CHAT_TEMPLATE_REGISTRY.register("Janus")

487class JanusTemplate(ChatTemplate):

488 """Render Janus multimodal conversations with image placeholder masks."""

489

490 def encode_messages(

491 self, messages: Sequence[Dict[str, str]], max_seq_len: int = 8192, task_type: str = ""

492 ) -> Dict[str, List[int]]:

493 """Encode messages into token, label, and image mask fields.

494

495 Args:

496 messages: Conversation records with ``role``, ``content``, and ``loss_mask`` fields.

497 max_seq_len: Maximum sequence length kept from the end of the encoding.

498 task_type: Optional task selector enabling generation-specific separators.

499

500 Returns:

501 Model inputs with input_ids, attention_mask, labels, images_seq_mask,

502 and images_emb_mask.

503

504 Raises:

505 ValueError: If a message role or image placeholder count is invalid.

506 """

507 # pylint: disable=too-many-locals

508 input_ids, attention_mask, labels = [], [], []

509 images_seq_mask, images_emb_mask = [], []

510 assistant_count = 0

511 for idx, message in enumerate(messages):

512 content_str, assistant_count = _format_janus_message(

513 message, idx, len(messages), task_type, assistant_count

514 )

515 if "system" in message["role"]:

516 content_ids = self.tokenizer.encode(content_str)

517 else:

518 content_ids = self.tokenizer.encode(content_str, add_special_tokens=False)

519 input_ids += content_ids

520 attention_mask += [1] * len(content_ids)

521 image_token_id = self.tokenizer.vocab.get("<image_placeholder>")

522 content_image_mask = [token == image_token_id for token in content_ids]

523 images_seq_mask += content_image_mask

524 num_image_tokens = sum(content_image_mask)

525 if num_image_tokens % 576:

526 raise ValueError("Each Janus image must have 576 placeholder tokens")

527 images_emb_mask.extend([[True] * 576 for _ in range(num_image_tokens // 576)])

528

529 labels += _janus_labels(content_ids, image_token_id, message["loss_mask"], task_type)

530

531 model_inputs = {

532 "input_ids": input_ids,

533 "attention_mask": attention_mask,

534 "labels": labels,

535 "images_seq_mask": images_seq_mask,

536 "images_emb_mask": images_emb_mask,

537 }

538 model_inputs = {k: v[-max_seq_len:] for k, v in model_inputs.items()}

539 return model_inputs

540

541 def get_jinja_template(self) -> str:

542 """Return the jinja template matching the Janus ChatML-style rendering."""

543 return (

544 "{% if not add_generation_prompt is defined %}{% set add_generation_prompt = false %}{% endif %}"

545 "{% for message in messages %}"

546 "{{ '<|im_start|>' + message['role'] + '\n' + message['content'] | trim + '<|im_end|>\n' }}"

547 "{% endfor %}"

548 "{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}"

549 )

550

551

552@CHAT_TEMPLATE_REGISTRY.register("chatml")

553class ChatmlTemplate(ChatTemplate):

554 """Render messages with the ChatML <|im_start|>/<|im_end|> format."""

555

556 def encode_messages(self, messages: Sequence[Dict[str, str]], max_seq_len: int = 8192) -> Dict[str, List[int]]:

557 """Encode messages into input_ids, attention_mask, and loss-masked labels.

558

559 Args:

560 messages: Conversation records with ``role``, ``content``, and optional ``loss_mask``.

561 max_seq_len: Maximum sequence length kept from the end of the encoding.

562

563 Returns:

564 Model inputs with input_ids, attention_mask, and labels.

565 """

566 input_ids, attention_mask, labels = [], [], []

567 for message in messages:

568 content_str = "<|im_start|>" + message["role"] + "\n" + message["content"].strip() + "<|im_end|>\n"

569 content_ids = self.tokenizer.encode(content_str, add_special_tokens=False)

570 input_ids += content_ids

571 attention_mask += [1] * len(content_ids)

572

573 if "loss_mask" in message:

574 loss_mask = message["loss_mask"]

575 else:

576 loss_mask = 1 if message["role"] == "assistant" else 0

577 if loss_mask == 1:

578 labels += content_ids

579 else:

580 labels += [IGNORE_INDEX] * len(content_ids)

581

582 model_inputs = {"input_ids": input_ids, "attention_mask": attention_mask, "labels": labels}

583 model_inputs = {k: v[-max_seq_len:] for k, v in model_inputs.items()}

584 return model_inputs

585

586 def get_jinja_template(self) -> str:

587 """Return the jinja template matching the ChatML rendering."""

588 return (

589 "{% if not add_generation_prompt is defined %}{% set add_generation_prompt = false %}{% endif %}"

590 "{% for message in messages %}"

591 "{{ '<|im_start|>' + message['role'] + '\n' + message['content'] | trim + '<|im_end|>\n' }}"

592 "{% endfor %}"

593 "{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}"

594 )

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

216 statements