File size: 4,728 Bytes
e3ac5f4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 | # Copyright (c) Gensyn Ltd.
# Licensed under the Apache License, Version 2.0.
"""Chat-aware tokenizer for open-1b.
Thin override of ``apply_chat_template`` that heals a BPE pretokenizer quirk:
encoding a string ending in whitespace (nothing after it) merges the trailing
run into different token(s) than when real content follows. The training
template's generation prompt is
``<|start_header|>assistant<|end_header|>\\n\\n`` — during SFT that ``\\n\\n``
is always training *context* immediately before the assistant's answer, never
a completion target (``sft.data.chat_template.render_conversation`` masks the
whole header, including the ``\\n\\n``, out of the loss). A plain encode of
the generation prompt on its own — nothing follows the ``\\n\\n`` — merges it
into a single token the model never saw at that position: its next-token
distribution reverts to the pretraining prior and it emits EOS immediately,
which looks fine on a single turn (the model usually recovers) but collapses
multi-turn chats to empty replies once a malformed turn enters the history.
We heal it the standard way: encode the prompt with a throwaway sentinel
character appended, then drop the sentinel's own token(s) — this forces real
content to follow the ``\\n\\n`` during encoding, so it tokenises the training
way (two separate tokens), and the sentinel's own tokens are cleanly
separable because it starts a new BPE pretoken. Same trick as
transformer-sft's ``sft.data.chat_template.encode_healed``, just applied
automatically so plain ``AutoTokenizer.apply_chat_template(...)`` users get
correct multi-turn behavior without knowing about the quirk.
"""
from __future__ import annotations
from transformers import PreTrainedTokenizerFast
from transformers.tokenization_utils_base import BatchEncoding
_HEAL_SENTINEL = "a"
class Open1BTokenizerFast(PreTrainedTokenizerFast):
def apply_chat_template(
self,
conversation,
tools=None,
documents=None,
chat_template=None,
add_generation_prompt=False,
continue_final_message=False,
tokenize=True,
padding=False,
truncation=False,
max_length=None,
return_tensors=None,
return_dict=True,
return_assistant_tokens_mask=False,
tokenizer_kwargs=None,
**kwargs,
):
# The quirk only bites the live-generation case: a single conversation,
# tokenized, with a bare (no-content-yet) generation prompt. Anything
# else — rendering text only, continuing/completing a final message,
# batched conversations, assistant-token masking for training — falls
# through to the standard implementation unchanged.
is_batch = conversation and isinstance(conversation[0], list)
if not (
tokenize and add_generation_prompt and not continue_final_message
and not return_assistant_tokens_mask and not is_batch
):
return super().apply_chat_template(
conversation,
tools=tools,
documents=documents,
chat_template=chat_template,
add_generation_prompt=add_generation_prompt,
continue_final_message=continue_final_message,
tokenize=tokenize,
padding=padding,
truncation=truncation,
max_length=max_length,
return_tensors=return_tensors,
return_dict=return_dict,
return_assistant_tokens_mask=return_assistant_tokens_mask,
tokenizer_kwargs=tokenizer_kwargs,
**kwargs,
)
text = super().apply_chat_template(
conversation,
tools=tools,
documents=documents,
chat_template=chat_template,
add_generation_prompt=True,
tokenize=False,
**kwargs,
)
ids = self._encode_healed(text)
encoded = BatchEncoding({"input_ids": [ids], "attention_mask": [[1] * len(ids)]})
if return_tensors is not None:
encoded.convert_to_tensors(tensor_type=return_tensors)
if not return_dict:
input_ids = encoded["input_ids"]
return input_ids[0] if return_tensors is None else input_ids
return encoded
def _encode_healed(self, text: str, sentinel: str = _HEAL_SENTINEL) -> list[int]:
full = self.encode(text + sentinel, add_special_tokens=False)
tail = self.encode(sentinel, add_special_tokens=False)
n = len(full) - len(tail)
if n >= 0 and full[n:] == tail:
return full[:n]
return self.encode(text, add_special_tokens=False)
|