human
4. Train the Model (Trainer)
4.1 Implement a train Python module.
Create a Python module containing a trainer_fn function, which must return an estimator.
The trainer_fn receives the following parameters:
* hparams: currently includes train_steps, eval_steps, and transform_output.
* schema: the raw data schema. This is used to create the serving input function to export the model.
The trainer_fn returns a dictionary of the following:
* estimator: The estimator that will be used for training and eval.
* train_spec: Spec for training.
* eval_spec: Spec for eval.
* eval_input_receiver_fn: Input function for eval.

Give me a the Python program you requested?
assistant
#Here is the Python program you requested.
_train_module_file = SRC_DIR + '/train.py'

%%writefile {_train_module_file}

import math

import tensorflow as tf
import tensorflow_transform as tft
import tensorflow_model_analysis as tfma
from tensorflow_transform.tf_metadata import schema_utils

TARGET_FEATURE_NAME = 'income_bracket'
WEIGHT_FEATURE_NAME = 'fnlwgt'
TARGET_FEATURE_LABELS = [' <=50K', ' >50K']

# Input function for train and eval transformed data
def make_input_fn(tfrecords_files, transform_output,
                  batch_size, num_epochs=1, shuffle=False):
    
    def _gzip_reader_fn(filenames):
        return tf.data.TFRecordDataset(
            filenames, compression_type='GZIP')
    
    def input_fn():
        dataset = tf.data.experimental.make_batched_features_dataset(
            file_pattern=tfrecords_files,
            batch_size=batch_size,
            features=transform_output.transformed_feature_spec(),
            label_key=TARGET_FEATURE_NAME,
            reader=_gzip_reader_fn,
            num_epochs=num_epochs,
            shuffle=shuffle
        )
        return dataset
    return input_fn

# Create estimator
def create_estimator(params, transform_output, run_config):
    
    feature_columns = []
    
    # Create feature columns based on the transform schema
    transformed_features = transform_output.transformed_metadata.schema.feature
    for feature in transformed_features:
        
        if feature.name in [TARGET_FEATURE_NAME, WEIGHT_FEATURE_NAME]:
            continue

        if hasattr(feature, 'int_domain') and feature.int_domain.is_categorical:
            vocab_size = feature.int_domain.max + 1
            feature_columns.append(
                tf.feature_column.embedding_column(
                    tf.feature_column.categorical_column_with_identity(
                        feature.name, num_buckets=vocab_size),
                    dimension = int(math.sqrt(vocab_size))))
        else:
            feature_columns.append(
                tf.feature_column.numeric_column(feature.name))
            
    #...