Class DataCollatorForCausalLMDecoderOnlyChatML

finetune/datacollators.py:14–242

This DataCollator is purpose-made for ChatML templates. Referenced: https://medium.com/@xuebinbin12/fine-tuning-chat-based-llm-with-multi-turn-conversational-data-part-i-d8c64d01a20d

@dataclass
class DataCollatorForCausalLMDecoderOnlyChatML(object):

    '''
    This DataCollator is purpose-made for ChatML templates.
    Referenced: https://medium.com/@xuebinbin12/fine-tuning-chat-based-llm-with-multi-turn-conversational-data-part-i-d8c64d01a20d
    '''
    tokenizer: PreTrainedTokenizer
    target_max_len: int

    # Decoder only doesn't require source source_max_length
    # source_max_length: int

    # `train_on_source` in this context refers to calculating prediction losses on the "user" prompts
    # In some cases where there are things to learn from the user prompts, this might help.
    # However in some finetuning dataset, the user prompts mimick a real-world poor quality prompts.
    # And usually doesn't contain anything useful. Therefore, prediction losses should be ignored,
    # So that the model doesn't learn to mimick it.
    train_on_source: bool  
    add_special_tokens: bool

    def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:

        if self.train_on_source:
            message = 'In some cases where there are things to learn from the user prompts, this might help.'
            message += 'However in some finetuning dataset, the user prompts mimick a real-world poor quality prompts.'
            message += 'And usually doesn\'t contain anything useful. Therefore, prediction losses should be ignored.'
            message += 'so that the model doesn\'t learn to mimick it.'
            match input(f'{message}\n\nAre you sure you want to continue? (yes/no): '):
                case 'yes':
                    pass
                case _:
                    exit()

        # In decoder-only model, we concat source and targets
        contents: List[str] = []

        for example in instances:
            content = example['input'] + example['output']
            contents.append(content)

        # Default is to not train on source, so we ignore losses on the source with IGNORE_INDEX
        if not self.train_on_source:

            # identify the start and end of the sequence to ignore losses
            # <|im_start|>user\nhello~<|im_end|>\n<|im_start|>assistant\nHow may I help?<|im_end|>
            # ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
            # The ^ represents the positions that we will apply IGNORE_INDEX

            # To do the ignore_index for cross-entropy loss, we need the start tokens and the end tokens to ignore
            start_tokens = self.tokenizer.encode(
                '<|im_start|>user\n',
                max_length=999,
                truncation=True,
                add_special_tokens=False,
            )
            end_tokens = self.tokenizer.encode(
                '<|im_start|>assistant\n',
                max_length=999,