Function process_fn

axlearn/common/input_lm.py:425–465

(example: dict[str, tf.Tensor])

    @seqio.utils.map_over_dataset
    def process_fn(example: dict[str, tf.Tensor]) -> dict[str, tf.Tensor]:
        target_labels = example["target_labels"]
        prefix = example["prefix"]  # This should match the "prefix" key used in decoding.
        input_ids = example["input_ids"]

        tf.debugging.assert_rank(target_labels, 1)
        tf.debugging.assert_rank(prefix, 1)
        tf.debugging.assert_rank(input_ids, 1)

        orig_target_labels_len = tf.size(target_labels)
        orig_input_ids_len = tf.size(input_ids)
        orig_prefix_len = tf.size(prefix)

        # Requires max_sequence_length to be greater than or equal to target_labels.
        tf.debugging.assert_less_equal(
            orig_target_labels_len,
            max_sequence_length,
            "max_sequence_length must be greater than target_labels length!",
        )
        total_tokens_to_truncate = tf.maximum(
            orig_input_ids_len + orig_target_labels_len + orig_prefix_len - 1 - max_sequence_length,
            0,
        )
        # Truncation starts from prefix.
        prefix_tokens_to_truncate = tf.minimum(total_tokens_to_truncate, orig_prefix_len)
        # Truncation continues on input_ids.
        input_ids_tokens_to_truncate = tf.minimum(
            total_tokens_to_truncate - prefix_tokens_to_truncate, orig_input_ids_len
        )
        # Truncate prefix from the left.
        prefix = prefix[prefix_tokens_to_truncate:]
        # Truncate input_ids from the left.
        input_ids = input_ids[input_ids_tokens_to_truncate:]
        # Under the above assertion, no target_labels should be truncated.
        return dict(
            prefix=prefix,
            target_labels=target_labels,
            input_ids=input_ids,
            prefix_tokens_to_truncate=prefix_tokens_to_truncate,
            input_ids_tokens_to_truncate=input_ids_tokens_to_truncate,
        )

    return process_fn