-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathdisaster_tweets_tuner.py
128 lines (107 loc) · 4.27 KB
/
disaster_tweets_tuner.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
import os
import tensorflow as tf
import tensorflow_transform as tft
import keras_tuner as kt
from tensorflow.keras import layers
from tfx.components.trainer.fn_args_utils import FnArgs
from keras_tuner.engine import base_tuner
from typing import NamedTuple, Dict, Text, Any
LABEL_KEY = 'label'
FEATURE_KEY = 'text'
def transformed_name(key):
"""Renaming transformed features"""
return key + "_xf"
def gzip_reader_fn(filenames):
"""Loads compressed data"""
return tf.data.TFRecordDataset(filenames, compression_type='GZIP')
def input_fn(file_pattern, tf_transform_output, num_epochs, batch_size=64) -> tf.data.Dataset:
"""Get post_tranform feature & create batches of data"""
# Get post_transform feature spec
transform_feature_spec = (
tf_transform_output.transformed_feature_spec().copy()
)
# create batches of data
dataset = tf.data.experimental.make_batched_features_dataset(
file_pattern = file_pattern,
batch_size = batch_size,
features = transform_feature_spec,
reader = gzip_reader_fn,
num_epochs = num_epochs,
label_key = transformed_name(LABEL_KEY)
)
return dataset
# Vocabulary size and number of words in a sequence.
VOCAB_SIZE = 10000
SEQUENCE_LENGTH = 100
vectorize_layer = layers.TextVectorization(
standardize = 'lower_and_strip_punctuation',
max_tokens = VOCAB_SIZE,
output_mode = 'int',
output_sequence_length = SEQUENCE_LENGTH
)
def model_builder(hp):
"""Build keras tuner model"""
embedding_dim = hp.Int('embedding_dim', min_value=16, max_value=128, step=16)
lstm_units = hp.Int('lstm_units', min_value=16, max_value=128, step=16)
num_layers = hp.Choice('num_layers', values=[1, 2, 3])
dense_units = hp.Int('dense_units', min_value=16, max_value=128, step=16)
dropout_rate = hp.Float('dropout_rate', min_value=0.1, max_value=0.5, step=0.1)
learning_rate = hp.Choice('learning_rate', values=[1e-2, 1e-3, 1e-4])
inputs = tf.keras.Input(shape=(1,), name=transformed_name(FEATURE_KEY), dtype=tf.string)
reshaped_narrative = tf.reshape(inputs, [-1])
x = vectorize_layer(reshaped_narrative)
x = layers.Embedding(VOCAB_SIZE, embedding_dim, name='embedding')(x)
x = layers.Bidirectional(layers.LSTM(lstm_units))(x)
for _ in range(num_layers):
x = layers.Dense(dense_units, activation='relu')(x)
x = layers.Dropout(dropout_rate)(x)
outputs = layers.Dense(1, activation='sigmoid')(x)
model = tf.keras.Model(inputs = inputs, outputs = outputs)
model.compile(
loss = tf.keras.losses.BinaryCrossentropy(from_logits=True),
optimizer = tf.keras.optimizers.Adam(learning_rate),
metrics = [tf.keras.metrics.BinaryAccuracy()]
)
model.summary()
return model
TunerFnResult = NamedTuple('TunerFnResult', [
('tuner', base_tuner.BaseTuner),
('fit_kwargs', Dict[Text, Any]),
])
early_stop_callback = tf.keras.callbacks.EarlyStopping(
monitor = 'val_binary_accuracy',
mode = 'max',
verbose = 1,
patience = 10
)
def tuner_fn(fn_args: FnArgs) -> None:
# Load the transform output
tf_transform_output = tft.TFTransformOutput(fn_args.transform_graph_path)
# Create batches of data
train_set = input_fn(fn_args.train_files[0], tf_transform_output, 10)
val_set = input_fn(fn_args.eval_files[0], tf_transform_output, 10)
vectorize_layer.adapt(
[j[0].numpy()[0] for j in [
i[0][transformed_name(FEATURE_KEY)]
for i in list(train_set)
]]
)
# Build the model tuner
model_tuner = kt.Hyperband(
hypermodel = lambda hp: model_builder(hp),
objective = kt.Objective('val_binary_accuracy', direction='max'),
max_epochs = 10,
factor = 3,
directory = fn_args.working_dir,
project_name = 'disaster_tweets_kt'
)
return TunerFnResult(
tuner = model_tuner,
fit_kwargs = {
'callbacks' : [early_stop_callback],
'x' : train_set,
'validation_data' : val_set,
'steps_per_epoch' : fn_args.train_steps,
'validation_steps' : fn_args.eval_steps
}
)