import h5py
import json
import gzip

layer_name_dict = {
    'Dense': 'denseLayer',
    'Dropout': 'dropoutLayer',
    'Embedding': 'embeddingLayer',
    'BatchNormalization': 'batchNormalizationLayer',
    'LSTM': 'rLSTMLayer',
    'GRU': 'rGRULayer',
    'JZS1': 'rJZS1Layer',
    'JZS2': 'rJZS2Layer',
    'JZS3': 'rJZS3Layer'
}

layer_params_dict = {
    'Dense': ['weights', 'activation'],
    'Dropout': ['p'],
    'Embedding': ['weights'],
    'BatchNormalization': ['weights', 'epsilon'],
    'LSTM': ['weights', 'activation', 'inner_activation'],
    'GRU': ['weights', 'activation', 'inner_activation'],
    'JZS1': ['weights', 'activation', 'inner_activation'],
    'JZS2': ['weights', 'activation', 'inner_activation'],
    'JZS3': ['weights', 'activation', 'inner_activation']
}

layer_weights_dict = {
    'Dense': ['W', 'b'],
    'Embedding': ['E'],
    'BatchNormalization': ['gamma', 'beta', 'mean', 'std'],
    'LSTM': ['W_xi', 'W_hi', 'b_i', 'W_xc', 'W_hc', 'b_c', 'W_xf', 'W_hf', 'b_f', 'W_xo', 'W_ho', 'b_o'],
    'GRU': ['W_xz', 'W_hz', 'b_z', 'W_xr', 'W_hr', 'b_r', 'W_xh', 'W_hh', 'b_h'],
    'JZS1': ['W_xz', 'b_z', 'W_xr', 'W_hr', 'b_r', 'W_hh', 'b_h', 'Pmat'],
    'JZS2': ['W_xz', 'W_hz', 'b_z', 'W_hr', 'b_r', 'W_xh', 'W_hh', 'b_h', 'Pmat'],
    'JZS3': ['W_xz', 'W_hz', 'b_z', 'W_xr', 'W_hr', 'b_r', 'W_xh', 'W_hh', 'b_h']
}

def serialize_from_model(model_json_file, weights_hdf5_file, save_filepath):
    with open(model_json_file, 'r') as f:
        model_metadata = json.load(f)
    weights_file = h5py.File(weights_hdf5_file, 'r')

    layers = []

    num_activation_layers = 0
    for k, layer in enumerate(model_metadata['layers']):
        if layer['name'] == 'Activation':
            num_activation_layers += 1
            prev_layer_name = layers[k-num_activation_layers]['layerName']
            idx_activation = layer_params_dict[prev_layer_name].index('activation')
            layers[k-num_activation_layers]['parameters'][idx_activation] = layer['activation']
            continue

        layer_params = []

        for param in layer_params_dict[layer['name']]:
            if param == 'weights':
                layer_weights = list(weights_file.get('layer_{}'.format(k)))
                weights = {}
                weight_names = layer_weights_dict[layer['name']]
                for name, w in zip(weight_names, layer_weights):
                    weights[name] = weights_file.get('layer_{}/{}'.format(k, w)).value.tolist()
                layer_params.append(weights)
            else:
                layer_params.append(layer[param])

        layers.append({
            'layerName': layer['name'],
            'parameters': layer_params
        })

    with gzip.open(save_filepath, 'wb') as f:
        f.write(json.dumps(layers).encode('utf8'))

@@ -1423,16 +1423,26 @@
 filepath
+, compress
 ):%0A    w
@@ -2710,16 +2710,37 @@
     %7D)%0A%0A
+    if compress:%0A    
     with
@@ -2785,16 +2785,20 @@
         
+    
 f.write(
@@ -2832,8 +2832,95 @@
 utf8'))%0A
+    else:%0A        with open(save_filepath, 'w') as f:%0A            json.dump(layers, f)%0A

update keras serialize script with option to apply gzip compression