| from __future__ import annotations | |
| import torch | |
| from experiments.train_qwen3vl_sft import _assistant_only_labels_and_weights | |
| class _CharacterTokenizer: | |
| im_end_id = 1 | |
| def __call__( | |
| self, | |
| text: str, | |
| *, | |
| add_special_tokens: bool = False, | |
| return_offsets_mapping: bool = False, | |
| ) -> dict[str, object]: | |
| del add_special_tokens | |
| value: dict[str, object] = { | |
| "input_ids": [ord(character) + 10 for character in text] | |
| } | |
| if return_offsets_mapping: | |
| value["offset_mapping"] = [ | |
| (index, index + 1) for index in range(len(text)) | |
| ] | |
| return value | |
| def convert_tokens_to_ids(self, token: str) -> int | None: | |
| return self.im_end_id if token == "<|im_end|>" else None | |
| class _BoundaryMergingTokenizer: | |
| def apply_chat_template( | |
| self, | |
| messages: list[dict], | |
| *, | |
| tokenize: bool, | |
| add_generation_prompt: bool, | |
| ) -> str: | |
| assert tokenize is False and add_generation_prompt is False | |
| return "".join( | |
| item["text"] | |
| for message in messages | |
| for item in message["content"] | |
| if item["type"] == "text" | |
| ) | |
| def __call__( | |
| self, | |
| text: str, | |
| *, | |
| add_special_tokens: bool = False, | |
| return_offsets_mapping: bool = False, | |
| ) -> dict[str, object]: | |
| assert add_special_tokens is False | |
| if text == "ABC": | |
| value: dict[str, object] = {"input_ids": [100, 101]} | |
| if return_offsets_mapping: | |
| value["offset_mapping"] = [(0, 2), (2, 3)] | |
| return value | |
| return {"input_ids": [200 + ord(character) for character in text]} | |
| def convert_tokens_to_ids(self, token: str) -> None: | |
| del token | |
| return None | |
| def test_only_assistant_semantic_answer_and_terminator_receive_loss() -> None: | |
| tokenizer = _CharacterTokenizer() | |
| question = "A" | |
| answer = "B" | |
| input_ids = torch.tensor( | |
| tokenizer(question)["input_ids"] | |
| + tokenizer(answer)["input_ids"] | |
| + [tokenizer.im_end_id], | |
| dtype=torch.long, | |
| ) | |
| messages = [ | |
| { | |
| "role": "user", | |
| "content": [{"type": "text", "text": question}], | |
| "loss": False, | |
| }, | |
| { | |
| "role": "assistant", | |
| "content": [{"type": "text", "text": answer}], | |
| "loss": True, | |
| }, | |
| ] | |
| labels, weights = _assistant_only_labels_and_weights( | |
| tokenizer, input_ids, messages | |
| ) | |
| assert labels.tolist() == [-100, ord("B") + 10, tokenizer.im_end_id] | |
| assert weights.tolist() == [0.0, 1.0, 1.0] | |
| def test_rendered_offsets_handle_tokens_crossing_message_boundaries() -> None: | |
| tokenizer = _BoundaryMergingTokenizer() | |
| messages = [ | |
| {"role": "user", "content": [{"type": "text", "text": "A"}]}, | |
| {"role": "assistant", "content": [{"type": "text", "text": "B"}]}, | |
| {"role": "user", "content": [{"type": "text", "text": "C"}]}, | |
| ] | |
| labels, weights = _assistant_only_labels_and_weights( | |
| tokenizer, torch.tensor([100, 101]), messages | |
| ) | |
| assert labels.tolist() == [100, -100] | |
| assert weights.tolist() == [1.0, 0.0] | |