Function preprocess_packed_supervised_dataset

src/utils/data_utils.py:136–175

(examples, tokenizer, cutoff_len)

    return knapsacks

def preprocess_packed_supervised_dataset(examples, tokenizer, cutoff_len):
    valid_num = 0
    batch_input_ids, batch_labels = [], []
    lengths = []
    length2indexes = defaultdict(list)
    for i in range(len(examples["input_ids"])):
        input_ids, labels = examples["input_ids"][i], examples["labels"][i]
        length = len(input_ids)
        if length >= cutoff_len - 1:
            continue
        else:
            lengths.append(length)
            length2indexes[length].append(valid_num)
            batch_input_ids.append(input_ids)
            batch_labels.append(labels)
            valid_num += 1
    model_inputs = defaultdict(list)
    knapsacks = greedy_knapsack(lengths, cutoff_len - 1)
    for knapsack in knapsacks:
        packed_input_ids, packed_attention_masks, packed_labels = [], [], []
        for i, length in enumerate(knapsack):
            index = length2indexes[length].pop()
            packed_input_ids += batch_input_ids[index]
            packed_labels += batch_labels[index]
            packed_attention_masks += [1] * len(batch_input_ids[index])

        if len(packed_input_ids) < cutoff_len:
            pad_length = cutoff_len - len(packed_input_ids)
            packed_input_ids += [tokenizer.pad_token_id] * pad_length
            packed_labels += [IGNORE_INDEX] * pad_length
            packed_attention_masks += [1] * pad_length  # more efficient flash_attn

        if len(packed_input_ids) != cutoff_len:
            raise ValueError("The length of packed example should be identical to the cutoff length.")

        model_inputs["input_ids"].append(packed_input_ids)
        model_inputs["attention_mask"].append(packed_attention_masks)
        model_inputs["position_ids"].append(list(range(len(packed_input_ids))))
        model_inputs["labels"].append(packed_labels)
    return model_inputs

def pad_sequence(examples, cutoff_len, tokenizer):
    max_length = cutoff_len