def tcn_featurizer(X, encoder, config, train=False, reuse=None, **kwargs):
"""
The featurizer element of the finetuning model. Maps from tokens ids to a dense embedding of the sequence.
:param X: A tensor of token indexes with shape [batch_size, sequence_length, token_idx]
:param encoder: A TextEncoder object.
:param config: A config object, containing all parameters for the featurizer.
:param train: If this flag is true, dropout and losses are added to the graph.
:param reuse: Should reuse be set within this scope.
:return: A dict containing;
embed_weights: the word embedding matrix.
features: The output of the featurizer_final state.
sequence_features: The output of the featurizer at each timestep.
"""
initial_shape = tf.shape(X)
X = tf.reshape(X, shape=tf.concat(([-1], initial_shape[-2:]), 0))
with tf.variable_scope("model/featurizer", reuse=reuse):
embed_weights = tf.get_variable(
name="we",
shape=[
encoder.vocab_size + config.max_length,
config.n_embed_featurizer
],
initializer=tf.random_normal_initializer(
stddev=config.weight_stddev),
)
if config.train_embeddings:
embed_weights = dropout(embed_weights, config.embed_p_drop, train)
else:
embed_weights = tf.stop_gradient(embed_weights)
X = tf.reshape(X, [-1, config.max_length, 2])
# we remove positional embeddings from the model
h = embed(X[:, :, :1], embed_weights)
# keep track of the classify token
clf_token = encoder["_classify_"]
with tf.variable_scope("tcn_stack"):
representation = h
for layer_num in range(config.n_layer):
representation = TemporalBlock(
n_filters=config.n_filter,
kernel_size=config.kernel_size,
rate=config.resid_p_drop if train else 0,
dilation_rate=2**layer_num,
scope="Temporal{}".format(layer_num),
)(representation)
seq_feats = tf.reshape(representation,
shape=[-1, config.max_length, config.n_filter])
# mask out the values past the classify token before performing pooling
pool_idx = tf.cast(
tf.argmax(tf.cast(tf.equal(X[:, :, 0], clf_token), tf.float32), 1),
tf.int32,
)
# mask is past the classify token (i.e. make those results extremely negative)
mask = tf.expand_dims(
1.0 - tf.sequence_mask(pool_idx,
maxlen=tf.shape(representation)[1],
dtype=tf.float32),
-1,
)
pool = tf.reduce_max(representation + mask * -1e9, 1)
clf_h = pool
clf_h = tf.reshape(clf_h,
shape=tf.concat(
(initial_shape[:-2], [config.n_filter]), 0))
# note that, due to convolution and pooling, the dimensionality of the features is much smaller than in the
# transformer base models
lengths = lengths_from_eos_idx(eos_idx=pool_idx,
max_length=config.max_length)
return {
"embed_weights": embed_weights,
"features":
clf_h, # [batch_size, n_embed] for classify, [batch_size, 1, n_embed] for comparison, etc.
"sequence_features": seq_feats, # [batch_size, seq_len, n_embed]
"eos_idx": pool_idx, # [batch_size]
"lengths": lengths
}
def featurizer(X, encoder, config, train=False, reuse=None, encoder_state=None, context=None, context_dim=None, **kwargs):
"""
The main element of the OSCAR model. Maps from tokens ids to a dense, embedding of the sequence.
:param X: A tensor of token indexes with shape [batch_size, sequence_length, token_idx]
:param encoder: A TextEncoder object.
:param config: A config object, containing all parameters for the featurizer.
:param train: If this flag is true, dropout and losses are added to the graph.
:param reuse: Should reuse be set within this scope.
:return: A dict containing;
embed_weights: the word embedding matrix.
features: The output of the featurizer_final state.
sequence_features: The output of the featurizer at each timestep.
"""
initial_shape = [a or -1 for a in X.get_shape().as_list()]
if len(initial_shape) != 3:
X = tf.reshape(X, shape=[-1] + initial_shape[-2:])
x_shape = tf.shape(X)
with tf.variable_scope('model/featurizer', reuse=reuse):
encoder._lazy_init()
clf_token = encoder.end_token
pool_idx = tf.cast(tf.argmax(tf.cast(tf.equal(X[:, :, 0], clf_token), tf.float32), 1), tf.int32)
if encoder_state is None:
embed_weights = tf.get_variable("we", [encoder.vocab_size + config.max_length, config.n_embed],
initializer=tf.random_normal_initializer(stddev=config.weight_stddev))
else:
embed_weights = encoder_state["embed_weights"]
if config.oscar_use_fp16:
embed_weights = tf.cast(embed_weights, tf.float16)
if config.train_embeddings:
embed_weights = dropout(embed_weights, config.embed_p_drop, train)
else:
embed_weights = tf.stop_gradient(embed_weights)
X = tf.reshape(X, [-1, x_shape[1], 2])
if config.oscar_use_timing:
h = embed(X, embed_weights)
else:
h = embed_no_timing(X, embed_weights)
for layer in range(config.n_layer):
with tf.variable_scope('h%d_' % layer):
if (
(config.n_layer - layer) == config.num_layers_trained and
config.num_layers_trained != config.n_layer
):
h = tf.stop_gradient(h)
block_fn_fwd = functools.partial(
block, block_name='block%d_' % layer, use_fp16=config.oscar_use_fp16,
pool_idx=None, encoder_state=encoder_state, train=train,
pdrop=config.resid_p_drop, use_fused_kernel=config.oscar_use_fused_kernel,
)
if config.low_memory_mode and train:
block_fn_fwd = recompute_grad(block_fn_fwd, use_entire_scope=True)
h = block_fn_fwd(h)
h = normal_1d_conv_block(h, 1, "output", config.oscar_use_fp16, dilation=1)
mask = tf.expand_dims(tf.sequence_mask(pool_idx, maxlen=tf.shape(h)[1], dtype=h.dtype), -1)
if config.oscar_feat_mode == "clf_tok":
clf_h = tf.gather_nd(h, tf.stack([tf.range(shape_list(h)[0]), pool_idx], 1))
elif config.oscar_feat_mode == "mean_tok":
clf_h = tf.reduce_sum(h * mask, 1) / tf.reduce_sum(h)
elif config.oscar_feat_mode == "max_tok":
clf_h = tf.reduce_max(h - (1e5 * (1.0 - mask)), 1)
else:
raise ValueError("config.feat_mode should be one of clf_tok, mean_tok or max_tok")
if len(initial_shape) != 3:
seq_feats = tf.reshape(h, shape=initial_shape[:-1] + [config.n_embed])
else:
seq_feats = h
return {
'embed_weights': embed_weights,
'features': cast_maybe(clf_h, tf.float32),
'sequence_features': seq_feats,
'eos_idx': pool_idx,
'encoded_input': X[:, :tf.reduce_min(pool_idx), 0],
'lengths': lengths_from_eos_idx(eos_idx=pool_idx, max_length=shape_list(X)[0])
}
def textcnn_featurizer(X, encoder, config, train=False, reuse=None):
"""
The transformer element of the finetuning model. Maps from tokens ids to a dense, embedding of the sequence.
:param X: A tensor of token indexes with shape [batch_size, sequence_length, token_idx]
:param encoder: A TextEncoder object.
:param config: A config object, containing all parameters for the featurizer.
:param train: If this flag is true, dropout and losses are added to the graph.
:param reuse: Should reuse be set within this scope.
:return: A dict containing;
embed_weights: the word embedding matrix.
features: The output of the featurizer_final state.
sequence_features: The output of the featurizer at each timestep.
"""
initial_shape = tf.shape(X)
X = tf.reshape(X, shape=tf.concat(([-1], initial_shape[-2:]), 0))
with tf.variable_scope('model/featurizer', reuse=reuse):
embed_weights = tf.get_variable(
name="we",
shape=[
encoder.vocab_size + config.max_length,
config.n_embed_featurizer
],
initializer=tf.random_normal_initializer(
stddev=config.weight_stddev))
if config.train_embeddings:
embed_weights = dropout(embed_weights, config.embed_p_drop, train)
else:
embed_weights = tf.stop_gradient(embed_weights)
X = tf.reshape(X, [-1, config.max_length, 2])
# we remove positional embeddings from the model
h = embed(X[:, :, :1], embed_weights)
# keep track of the classify token
clf_token = encoder['_classify_']
# Convolutional Layer (this is all the same layer, just different filter sizes)
pool_layers = []
conv_layers = []
for i, kernel_size in enumerate(config.kernel_sizes):
conv = tf.layers.conv1d(
inputs=h,
filters=config.num_filters_per_size,
kernel_size=kernel_size,
padding='same',
activation=tf.nn.relu,
name='conv' + str(i),
kernel_initializer=tf.initializers.glorot_normal)
conv_layers.append(conv)
# mask out the values past the classify token before performing pooling
pool_idx = tf.cast(
tf.argmax(tf.cast(tf.equal(X[:, :, 0], clf_token), tf.float32),
1), tf.int32)
# mask is past the classify token (i.e. make those results extremely negative)
mask = tf.expand_dims(
1.0 - tf.sequence_mask(
pool_idx, maxlen=tf.shape(conv)[1], dtype=tf.float32), -1)
pool = tf.reduce_max(conv + mask * -1e9, 1)
pool_layers.append(pool)
# Concat the output of the convolutional layers for use in sequence embedding
conv_seq = tf.concat(conv_layers, axis=2)
seq_feats = tf.reshape(conv_seq,
shape=[-1, config.max_length, config.n_embed])
# Concatenate the univariate vectors as features for classification
clf_h = tf.concat(pool_layers, axis=1)
clf_h = tf.reshape(clf_h,
shape=tf.concat(
(initial_shape[:-2], [config.n_embed]), 0))
# note that, due to convolution and pooling, the dimensionality of the features is much smaller than in the
# transformer base models
return {
'embed_weights': embed_weights,
'features':
clf_h, # [batch_size, n_embed] for classify, [batch_size, 1, n_embed] for comparison, etc.
'sequence_features': seq_feats, # [batch_size, seq_len, n_embed]
'pool_idx': pool_idx # [batch_size]
}