Function rekey

axlearn/common/input_tf_data.py:1087–1142

Replace the feature keys according to mapping in `key_map`. Like seqio's rekey, except: 1. We allow for a configurable default value (used if the reference key is falsey--e.g. None--or if missing in the input example). 2. We optionally allow retaining keys not explic

(
    key_map: dict[str, str],
    default_value: Optional[Any] = "",
    retain_original_inputs: bool = False,
    separator: Optional[str] = None,
)

def rekey(
    key_map: dict[str, str],
    default_value: Optional[Any] = "",
    retain_original_inputs: bool = False,
    separator: Optional[str] = None,
) -> DatasetToDatasetFn:
    """Replace the feature keys according to mapping in `key_map`.

    Like seqio's rekey, except:
        1. We allow for a configurable default value
            (used if the reference key is falsey--e.g. None--or if missing in the input example).
        2. We optionally allow retaining keys not explicitly mentioned in the key-map.
        3. We optionally allow keys to be paths (if separator is provided).

    Ref: <https://github.com/google/seqio/blob/9748501b/seqio/preprocessors.py#L30-L52>

    Args:
        key_map: A dictionary mapping new keys to original keys.
            If falsey, return input (to match seqio behavior).
        default_value: Value to set new key to if old key-value doesn&#x27;t exist.
            If None, then we do not write the new key-value pair when missing an old key-value
                or when the provided reference key is falsey (to match seqio).
        retain_original_inputs: Whether to retain all the keys provided in the original input
            example (if False, only keys specified in the key map will be in the output).
        separator: An optional separator. If provided, all keys and values of `key_map` will be
            treated as paths and split by the separator.

    Returns:
        A DatasetToDatasetFn, where each input example should be a dict.
    """

    def has_path(x, path: str) -> bool:
        try:
            get_recursively(x, path, separator=separator)
            return True
        except KeyError:
            return False

    def fn(example: dict[str, tf.Tensor]) -> dict[str, tf.Tensor]:
        if not key_map:
            return example
        output = example if retain_original_inputs else {}
        for new_key, old_key in key_map.items():
            if not old_key or not has_path(example, old_key):
                if default_value is not None:
                    set_recursively(output, value=default_value, path=new_key, separator=separator)
                continue
            set_recursively(
                output,
                value=get_recursively(example, old_key, separator=separator),
                path=new_key,
                separator=separator,
            )
        return output

    return seqio.map_over_dataset(fn)