transform.py

LABEL_KEY = 'lbl'
DENSE_FEATURES = ['a', 'b']
BINARY_FEATURES = ['c']

def preprocessing_fn(inputs: Dict[Text, Any]) -> Dict[Text, Any]:
    outputs = {}
    for feat in DENSE_FEATURES:
        outputs[f'{feat}_xf'] = tft.scale_to_z_score(inputs[feat])

    for feat in BINARY_FEATURES:
        outputs[feat] = inputs[feat]

    outputs[LABEL_KEY] = inputs[LABEL_KEY]
    return outputs