model_utils.py

from typing import Tuple, Text, Dict
import tensorflow as tf
from tensorflow.keras.layers import Input, Concatenate, Reshape, Dense
from input_fn_utils import transformed_name


def get_input_graph(input_feature_keys, input_window_size) -> Tuple[Input, tf.keras.layers.Layer]:
    """ 
        Creates the named input layers, strips the column names and provides
        them as a plain tensor.
    Returns:
        Tuple[Input, tf.keras.layers.Layer]: Input for your model and a layer with output shape
        [None (batch size), input_window_size, len(input_feature_keys)]
    """
    transformed_columns = [transformed_name(
        key) for key in input_feature_keys]

    input_layers = {
        colname: Input(name=colname, shape=(
            input_window_size), dtype=tf.float32)
        for colname in transformed_columns
    }

    pre_model_input = Concatenate(axis=-1)(list(input_layers.values()))
    pre_model_input = Reshape(target_shape=(input_window_size, len(input_feature_keys)))(
        pre_model_input)

    return input_layers, pre_model_input


def get_output_graph(head_layer, predict_feature_keys, output_window_size) -> Dict[Text, tf.keras.layers.Layer]:
    """
        Transforms a plain-tensor feature layer output to named output layers.
    Args:
        head_layer ([type]): The final feature layer of your model
    Returns:
        Dict[Text, tf.keras.layers.Layer]: Named Dense layer outputs based on predict_feature_keys
    """
    return {
        colname: Dense(units=output_window_size, name=colname)(head_layer)
        for colname in predict_feature_keys
    }