From 32d5d8b5f67e2e50a7e5aa7c485b0bce4493ffac Mon Sep 17 00:00:00 2001 From: brainsqueeze Date: Wed, 6 Jul 2022 08:17:29 -0400 Subject: [PATCH 01/10] refactoring, better decode attn --- text2vec/models/components/attention.py | 68 +++++++++----------- text2vec/models/components/feed_forward.py | 9 +-- text2vec/models/components/feeder.py | 36 +++++------ text2vec/models/components/recurrent.py | 36 +++++------ text2vec/models/components/strings.py | 5 +- text2vec/models/components/text_inputs.py | 25 ++++---- text2vec/models/components/utils.py | 40 ++++++------ text2vec/models/sequential.py | 38 +++++------ text2vec/models/transformer.py | 74 ++++++++++------------ 9 files changed, 160 insertions(+), 171 deletions(-) diff --git a/text2vec/models/components/attention.py b/text2vec/models/components/attention.py index 07cbfa3..bfdf224 100644 --- a/text2vec/models/components/attention.py +++ b/text2vec/models/components/attention.py @@ -1,7 +1,10 @@ +from typing import Optional + import tensorflow as tf +from tensorflow.keras import layers, initializers -class ScaledDotAttention(tf.keras.layers.Layer): +class ScaledDotAttention(layers.Layer): """Scaled dot attention layer which computes ``` softmax(Query * permutedim(Key, (3, 1, 2)) / sqrt(dk)) * permutedim(Value, (2, 1, 3)) @@ -50,7 +53,7 @@ def call(self, query, key, value, mask_future=False): return tf.einsum('ijk,ikl->ijl', x, value) -class BahdanauAttention(tf.keras.layers.Layer): +class BahdanauAttention(layers.Layer): """Layer which computes the Bahdanau attention mechanism either as a self-attention or as a encoder-decoder attention. @@ -82,38 +85,29 @@ class BahdanauAttention(tf.keras.layers.Layer): ``` """ - def __init__(self, size): + def __init__(self, size: int, dropout_rate: float = 0.): super().__init__(name="BahdanauAttention") - initializer = tf.keras.initializers.GlorotUniform() - self.W = tf.Variable( - initializer(shape=(size, size)), - name='weight', - dtype=tf.float32, - trainable=True - ) - self.B = tf.Variable(tf.zeros(shape=[size]), name="B", dtype=tf.float32, trainable=True) - self.U = tf.Variable(initializer(shape=[size]), name="U", dtype=tf.float32, trainable=True) - - def call(self, encoded, decoded=None): - with tf.name_scope("BahdanauAttention"): - if decoded is None: - score = tf.math.tanh(tf.tensordot(encoded, self.W, axes=[-1, 0]) + self.B) - score = tf.reduce_sum(self.U * score, axis=-1) - alphas = tf.nn.softmax(score, name="attention-weights") - # encoded = encoded * tf.expand_dims(alphas, axis=-1) - # return encoded, tf.reduce_sum(encoded, axis=1, name="context-vector") - return tf.einsum('ilk,il->ik', encoded, alphas) - - score = tf.einsum("ijm,mn,ikn->ijk", encoded, self.W, decoded) - alphas = tf.reduce_mean(score, axis=1) - alphas = tf.nn.softmax(alphas) - # decoded = decoded * tf.expand_dims(alphas, axis=-1) - # return decoded, tf.reduce_sum(decoded, axis=1) - return tf.einsum('ilk,il->ik', decoded, alphas) - - -class SingleHeadAttention(tf.keras.layers.Layer): + self.hidden = layers.Dense(units=size, activation="tanh") + self.U = tf.Variable(initializers.GlorotUniform()(shape=[size]), name="U", dtype=tf.float32, trainable=True) + self.dropout = layers.Dropout(dropout_rate) + + def call(self, encoded: tf.Tensor, decoded: Optional[tf.Tensor] = None, training: bool = False) -> tf.Tensor: + if decoded is None: + score = tf.math.reduce_sum(self.U * self.hidden(encoded), axis=-1) + alphas = tf.nn.softmax(score) + alphas = self.dropout(alphas, training=training) + x = tf.expand_dims(alphas, axis=-1) * encoded + return x, tf.math.reduce_sum(x, axis=1) + # return tf.einsum('ilk,il->ik', encoded, alphas) + + score = tf.einsum("ijm,mn,ikn->ijk", encoded, self.hidden.kernel, decoded) + alphas = tf.nn.softmax(score, axis=1) + alphas = tf.math.reduce_sum(tf.matmul(alphas, encoded, transpose_a=True), axis=-1) + return tf.einsum('ilk,il->ik', decoded, alphas) + + +class SingleHeadAttention(layers.Layer): """Layer which computes the single-head-attention mechanism as described in https://arxiv.org/abs/1706.03762. @@ -160,10 +154,10 @@ def __init__(self, emb_dims, layers=8, keep_prob=1.0): self.WQ = tf.Variable(initializer(shape=(dims, key_dims)), name="WQ", dtype=tf.float32, trainable=True) self.WK = tf.Variable(initializer(shape=(dims, key_dims)), name="WK", dtype=tf.float32, trainable=True) self.WV = tf.Variable(initializer(shape=(dims, key_dims)), name="WV", dtype=tf.float32, trainable=True) - self.dropout = tf.keras.layers.Dropout(1 - keep_prob) + self.dropout = layers.Dropout(1 - keep_prob) self.dot_attention = ScaledDotAttention() - def call(self, inputs, mask_future=False, training=False): + def call(self, inputs, mask_future: bool = False, training: bool = False): with tf.name_scope("SingleHeadAttention"): queries, keys, values = inputs @@ -177,7 +171,7 @@ def call(self, inputs, mask_future=False, training=False): return self.dot_attention(query=head_queries, key=head_keys, value=head_values, mask_future=mask_future) -class MultiHeadAttention(tf.keras.layers.Layer): +class MultiHeadAttention(layers.Layer): """Layer which computes the multi-head-attention mechanism as described in https://arxiv.org/abs/1706.03762. @@ -213,14 +207,14 @@ class MultiHeadAttention(tf.keras.layers.Layer): ``` """ - def __init__(self, emb_dims, layers=8, keep_prob=1.0): + def __init__(self, emb_dims: int, layers: int = 8, keep_prob: float = 1.0): super().__init__(name="MultiHeadAttention") self.layer_heads = [] for i in range(layers): with tf.name_scope(f"head-{i}"): self.layer_heads.append(SingleHeadAttention(emb_dims=emb_dims, layers=layers, keep_prob=keep_prob)) - self.dense = tf.keras.layers.Dense(units=emb_dims, use_bias=False) + self.dense =layers.Dense(units=emb_dims, use_bias=False) def call(self, inputs, mask_future=False, training=False): with tf.name_scope("MultiHeadAttention"): diff --git a/text2vec/models/components/feed_forward.py b/text2vec/models/components/feed_forward.py index 012c829..ec80097 100644 --- a/text2vec/models/components/feed_forward.py +++ b/text2vec/models/components/feed_forward.py @@ -1,7 +1,8 @@ import tensorflow as tf +from tensorflow.keras import layers -class PositionWiseFFN(tf.keras.layers.Layer): +class PositionWiseFFN(layers.Layer): """Position-wise feed-forward network implemented as conv -> relu -> conv. 1D convolutions of the input tensor are computed to an intermediate hidden dimension, then a final 1D convolution is computed of the ReLu output from the intermediate layer to return to the original input shape. @@ -25,17 +26,17 @@ class PositionWiseFFN(tf.keras.layers.Layer): ``` """ - def __init__(self, emb_dims): + def __init__(self, emb_dims: int): super().__init__() - self.conv_inner = tf.keras.layers.Conv1D( + self.conv_inner = layers.Conv1D( filters=4 * emb_dims, kernel_size=1, padding='same', use_bias=False, activation='relu' ) - self.conv_outer = tf.keras.layers.Conv1D(filters=emb_dims, kernel_size=1, padding='same', use_bias=False) + self.conv_outer = layers.Conv1D(filters=emb_dims, kernel_size=1, padding='same', use_bias=False) def call(self, x): with tf.name_scope("PositionWiseFFN"): diff --git a/text2vec/models/components/feeder.py b/text2vec/models/components/feeder.py index b617ff6..5ad773e 100644 --- a/text2vec/models/components/feeder.py +++ b/text2vec/models/components/feeder.py @@ -1,7 +1,8 @@ import tensorflow as tf +from tensorflow.keras import layers -class Tokenizer(tf.keras.layers.Layer): +class Tokenizer(layers.Layer): """String-splitting layer. Parameters @@ -24,7 +25,7 @@ class Tokenizer(tf.keras.layers.Layer): ``` """ - def __init__(self, sep=' '): + def __init__(self, sep: str = ' '): super().__init__(name="Tokenizer") self.sep = sep @@ -32,7 +33,7 @@ def call(self, corpus): return tf.strings.split(corpus, self.sep) -class TextInput(tf.keras.layers.Layer): +class TextInput(layers.Layer): """This layer handles the primary text feature transformations and word-embeddings to be passed off to the sequence-aware parts of the encoder/decoder pipeline. @@ -75,7 +76,7 @@ class TextInput(tf.keras.layers.Layer): ``` """ - def __init__(self, token_hash, embedding_size, max_sequence_len): + def __init__(self, token_hash: dict, embedding_size: int, max_sequence_len: int): super().__init__() assert isinstance(token_hash, dict) @@ -95,21 +96,20 @@ def __init__(self, token_hash, embedding_size, max_sequence_len): trainable=True ) self.max_len = tf.constant(max_sequence_len) - self.slicer = tf.keras.layers.Lambda(lambda x: x[:, :max_sequence_len], name="sequence-slice") + self.slicer = layers.Lambda(lambda x: x[:, :max_sequence_len], name="sequence-slice") - def call(self, tokens, output_embeddings=False): - with tf.name_scope("TextInput"): - hashed = tf.ragged.map_flat_values(self.table.lookup, tokens) - hashed = self.slicer(hashed) + def call(self, tokens: tf.RaggedTensor, output_embeddings: bool = False): + hashed = tf.ragged.map_flat_values(self.table.lookup, tokens) + hashed = self.slicer(hashed) - x = tf.ragged.map_flat_values(tf.nn.embedding_lookup, self.embeddings, hashed) - if output_embeddings: - return x + x = tf.ragged.map_flat_values(tf.nn.embedding_lookup, self.embeddings, hashed) + if output_embeddings: + return x - x = x.to_tensor(0) - x = x * tf.math.sqrt(tf.cast(tf.shape(self.embeddings)[-1], tf.float32)) # sqrt(embedding_size) + x = x.to_tensor(0) + x = x * tf.math.sqrt(tf.cast(tf.shape(self.embeddings)[-1], tf.float32)) # sqrt(embedding_size) - seq_lengths = hashed.row_lengths() - time_steps = tf.cast(tf.reduce_max(seq_lengths), tf.int32) - mask = tf.sequence_mask(lengths=seq_lengths, maxlen=time_steps, dtype=tf.float32) - return x, mask, time_steps + seq_lengths = hashed.row_lengths() + time_steps = tf.cast(tf.reduce_max(seq_lengths), tf.int32) + mask = tf.sequence_mask(lengths=seq_lengths, maxlen=time_steps, dtype=tf.float32) + return x, mask, time_steps diff --git a/text2vec/models/components/recurrent.py b/text2vec/models/components/recurrent.py index be7d272..b2c18ed 100644 --- a/text2vec/models/components/recurrent.py +++ b/text2vec/models/components/recurrent.py @@ -1,7 +1,8 @@ import tensorflow as tf +from tensorflow.keras import layers -class BidirectionalLSTM(tf.keras.layers.Layer): +class BidirectionalLSTM(layers.Layer): """Bi-directional LSTM with the option to warm initialize with previous states. Parameters @@ -30,11 +31,11 @@ class BidirectionalLSTM(tf.keras.layers.Layer): ``` """ - def __init__(self, num_layers=2, num_hidden=32, return_states=False): + def __init__(self, num_layers: int = 2, num_hidden: int = 32, return_states: bool = False): super().__init__() self.num_layers = num_layers self.return_states = return_states - lstm = tf.keras.layers.LSTM + lstm = layers.LSTM params = dict( units=num_hidden, @@ -45,7 +46,7 @@ def __init__(self, num_layers=2, num_hidden=32, return_states=False): self.FWD = [lstm(**params, name=f"forward-{i}") for i in range(num_layers)] self.BWD = [lstm(**params, name=f"backward-{i}", go_backwards=True) for i in range(num_layers)] - self.concat = tf.keras.layers.Concatenate() + self.concat = layers.Concatenate() @staticmethod def __make_inputs(inputs, initial_states=None, layer=0): @@ -59,19 +60,18 @@ def __make_inputs(inputs, initial_states=None, layer=0): return fwd_inputs, bwd_inputs def call(self, inputs, initial_states=None, training=False): - with tf.name_scope("BidirectionalLSTM"): - layer = 0 - for forward, backward in zip(self.FWD, self.BWD): - fwd_inputs, bwd_inputs = self.__make_inputs(inputs, initial_states=initial_states, layer=layer) + layer = 0 + for forward, backward in zip(self.FWD, self.BWD): + fwd_inputs, bwd_inputs = self.__make_inputs(inputs, initial_states=initial_states, layer=layer) - if self.return_states: - decode_forward, *forward_state = forward(**fwd_inputs, training=training) - decode_backward, *backward_state = backward(**bwd_inputs, training=training) - else: - decode_forward = forward(**fwd_inputs, training=training) - decode_backward = backward(**bwd_inputs, training=training) - inputs = self.concat([decode_forward, decode_backward]) - layer += 1 if self.return_states: - return inputs, [forward_state, backward_state] - return inputs + decode_forward, *forward_state = forward(**fwd_inputs, training=training) + decode_backward, *backward_state = backward(**bwd_inputs, training=training) + else: + decode_forward = forward(**fwd_inputs, training=training) + decode_backward = backward(**bwd_inputs, training=training) + inputs = self.concat([decode_forward, decode_backward]) + layer += 1 + if self.return_states: + return inputs, [forward_state, backward_state] + return inputs diff --git a/text2vec/models/components/strings.py b/text2vec/models/components/strings.py index 573628e..b8781f8 100644 --- a/text2vec/models/components/strings.py +++ b/text2vec/models/components/strings.py @@ -1,9 +1,10 @@ import tensorflow as tf +from tensorflow.kersa import layers from text2vec.models import Tokenizer -class SubStringFinderMask(tf.keras.layers.Layer): +class SubStringFinderMask(layers.Layer): """Performs substring masking based on whether the substring is found in the input text in its entirety. This returns a ragged boolean tensor with the same ragged shape as input substrings. @@ -35,7 +36,7 @@ class SubStringFinderMask(tf.keras.layers.Layer): def __init__(self, sep: str = ' '): super().__init__() self.tokenizer = Tokenizer(sep) - self.match = tf.keras.layers.Lambda(lambda x: tf.strings.regex_full_match(input=x[0], pattern=x[1])) + self.match = layers.Lambda(lambda x: tf.strings.regex_full_match(input=x[0], pattern=x[1])) # this is designed to approximate the functionality in re.escape self.special_chars = r'[\(\)\[\]\{\}\?\*\+\-\|\^\$\\\\\.\&\~\#\\\t\\\n\\\r\\\v\\\f]' diff --git a/text2vec/models/components/text_inputs.py b/text2vec/models/components/text_inputs.py index 062aea9..888ef4b 100644 --- a/text2vec/models/components/text_inputs.py +++ b/text2vec/models/components/text_inputs.py @@ -1,7 +1,8 @@ import tensorflow as tf +from tensorflow.keras import layers -class Embed(tf.keras.layers.Layer): +class Embed(layers.Layer): """This layer handles the primary text feature transformations and word-embeddings to be passed off to the sequence-aware parts of the encoder/decoder pipeline. @@ -48,19 +49,18 @@ def __init__(self, vocab_size: int, embedding_size: int, max_sequence_len: int): trainable=True ) self.max_len = tf.constant(max_sequence_len) - self.slicer = tf.keras.layers.Lambda(lambda x: x[:, :max_sequence_len], name="sequence-slice") + self.slicer = layers.Lambda(lambda x: x[:, :max_sequence_len], name="sequence-slice") def call(self, token_ids, **kwargs): - with tf.name_scope("TokenIds"): - token_ids = self.slicer(token_ids) - x = tf.ragged.map_flat_values(tf.nn.embedding_lookup, self.embeddings, token_ids) - x = x.to_tensor(0) - x = x * tf.math.sqrt(tf.cast(tf.shape(self.embeddings)[-1], tf.float32)) # sqrt(embedding_size) + token_ids = self.slicer(token_ids) + x = tf.ragged.map_flat_values(tf.nn.embedding_lookup, self.embeddings, token_ids) + x = x.to_tensor(0) + x = x * tf.math.sqrt(tf.cast(tf.shape(self.embeddings)[-1], tf.float32)) # sqrt(embedding_size) - seq_lengths = token_ids.row_lengths() - time_steps = tf.cast(tf.reduce_max(seq_lengths), tf.int32) - mask = tf.sequence_mask(lengths=seq_lengths, maxlen=time_steps, dtype=tf.float32) - return x, mask, time_steps + seq_lengths = token_ids.row_lengths() + time_steps = tf.cast(tf.reduce_max(seq_lengths), tf.int32) + mask = tf.sequence_mask(lengths=seq_lengths, maxlen=time_steps, dtype=tf.float32) + return x, mask, time_steps def get_embedding(self, token_ids: tf.RaggedTensor) -> tf.RaggedTensor: """Get the token embeddings for the input IDs. @@ -76,8 +76,7 @@ def get_embedding(self, token_ids: tf.RaggedTensor) -> tf.RaggedTensor: Sequences of token embeddings with the same number of time steps as `token_ids` """ - with tf.name_scope("TokenEmbeddings"): - return tf.ragged.map_flat_values(tf.nn.embedding_lookup, self.embeddings, token_ids) + return tf.ragged.map_flat_values(tf.nn.embedding_lookup, self.embeddings, token_ids) class TokenEmbed(tf.keras.layers.Layer): diff --git a/text2vec/models/components/utils.py b/text2vec/models/components/utils.py index dd79ee6..9da7012 100644 --- a/text2vec/models/components/utils.py +++ b/text2vec/models/components/utils.py @@ -1,8 +1,9 @@ import numpy as np import tensorflow as tf +from tensorflow.keras import layers -class LayerNorm(tf.keras.layers.Layer): +class LayerNorm(layers.Layer): """Layer normalization, independent of batch size. Parameters @@ -26,21 +27,20 @@ class LayerNorm(tf.keras.layers.Layer): ``` """ - def __init__(self, epsilon=1e-8, scale=1.0, bias=0): + def __init__(self, epsilon: float = 1e-8, scale: float = 1.0, bias: float = 0): super().__init__(name="LayerNorm") self.epsilon = tf.constant(epsilon, dtype=tf.float32) self.scale = tf.constant(scale, dtype=tf.float32) self.bias = tf.constant(bias, dtype=tf.float32) def call(self, x): - with tf.name_scope("LayerNorm"): - mean = tf.reduce_mean(x, axis=-1, keepdims=True) - variance = tf.reduce_mean(tf.square(x - mean), axis=-1, keepdims=True) - norm = (x - mean) * tf.math.rsqrt(variance + self.epsilon) - return norm * self.scale + self.bias + mean = tf.reduce_mean(x, axis=-1, keepdims=True) + variance = tf.reduce_mean(tf.square(x - mean), axis=-1, keepdims=True) + norm = (x - mean) * tf.math.rsqrt(variance + self.epsilon) + return norm * self.scale + self.bias -class TensorProjection(tf.keras.layers.Layer): +class TensorProjection(layers.Layer): """Projects sequence vectors onto a fixed vector. This returns a new tensor with the same shape as the input tensor, with all sequence vectors projected. @@ -63,17 +63,16 @@ def __init__(self): super().__init__(name="TensorProjection") def call(self, x, projection_vector): - with tf.name_scope("TensorProjection"): - inner_product = tf.einsum("ijk,ik->ij", x, projection_vector) - time_steps = tf.shape(x)[1] - p_vector_norm_squared = tf.norm(projection_vector, axis=1) ** 2 - p_vector_norm_squared = tf.tile(tf.expand_dims(p_vector_norm_squared, -1), [1, time_steps]) + inner_product = tf.einsum("ijk,ik->ij", x, projection_vector) + time_steps = tf.shape(x)[1] + p_vector_norm_squared = tf.norm(projection_vector, axis=1) ** 2 + p_vector_norm_squared = tf.tile(tf.expand_dims(p_vector_norm_squared, -1), [1, time_steps]) - alpha = tf.divide(inner_product, p_vector_norm_squared) - return tf.einsum("ij,ik->ijk", alpha, projection_vector) + alpha = tf.divide(inner_product, p_vector_norm_squared) + return tf.einsum("ij,ik->ijk", alpha, projection_vector) -class PositionalEncoder(tf.keras.layers.Layer): +class PositionalEncoder(layers.Layer): """Layer which initializes the positional encoding tensor, and defines the operation which adds the encoding to an input tensor and then applies a sequence mask. @@ -104,7 +103,7 @@ class PositionalEncoder(tf.keras.layers.Layer): ``` """ - def __init__(self, emb_dims, max_sequence_len): + def __init__(self, emb_dims: int, max_sequence_len: int): super().__init__() positions = np.arange(max_sequence_len).astype(np.float32) @@ -119,7 +118,6 @@ def __init__(self, emb_dims, max_sequence_len): encoder[:, 1::2] = odd self.encoder = tf.convert_to_tensor(encoder, dtype=tf.float32) - def call(self, x, mask): - with tf.name_scope('PositionalEncoder'): - time_steps = tf.shape(x)[1] - return tf.einsum('ijk,ij->ijk', x + self.encoder[:time_steps, :], mask) + def call(self, x: tf.Tensor, mask: tf.Tensor): + time_steps = tf.shape(x)[1] + return tf.einsum('ijk,ij->ijk', x + self.encoder[:time_steps, :], mask) diff --git a/text2vec/models/sequential.py b/text2vec/models/sequential.py index c740cf5..46a5869 100644 --- a/text2vec/models/sequential.py +++ b/text2vec/models/sequential.py @@ -1,11 +1,12 @@ import tensorflow as tf +from tensorflow.keras import layers from .components.attention import BahdanauAttention from .components.recurrent import BidirectionalLSTM from .components.utils import TensorProjection -class RecurrentEncoder(tf.keras.layers.Layer): +class RecurrentEncoder(layers.Layer): """LSTM based encoding pipeline. Parameters @@ -43,7 +44,7 @@ def __init__(self, max_sequence_len, num_hidden, num_layers=2, input_keep_prob=1 super().__init__() self.max_sequence_length = max_sequence_len - self.drop = tf.keras.layers.Dropout(1 - input_keep_prob, name="InputDropout") + self.drop = layers.Dropout(1 - input_keep_prob, name="InputDropout") self.bi_lstm = BidirectionalLSTM(num_layers=num_layers, num_hidden=num_hidden, return_states=True) self.attention = BahdanauAttention(size=2 * num_hidden) @@ -59,7 +60,7 @@ def call(self, x, mask, training=False, **kwargs): return x, context -class RecurrentDecoder(tf.keras.layers.Layer): +class RecurrentDecoder(layers.Layer): """LSTM based decoding pipeline. Parameters @@ -84,24 +85,23 @@ def __init__(self, max_sequence_len, num_hidden, embedding_size=50, num_layers=2 self.max_sequence_length = max_sequence_len dims = embedding_size - self.drop = tf.keras.layers.Dropout(1 - input_keep_prob, name="InputDropout") - self.h_drop = tf.keras.layers.Dropout(1 - hidden_keep_prob, name="HiddenStateDropout") + self.drop = layers.Dropout(1 - input_keep_prob, name="InputDropout") + self.h_drop = layers.Dropout(1 - hidden_keep_prob, name="HiddenStateDropout") self.projection = TensorProjection() self.bi_lstm = BidirectionalLSTM(num_layers=num_layers, num_hidden=num_hidden, return_states=False) - self.dense = tf.keras.layers.Dense(units=dims, activation=tf.nn.relu) + self.dense = layers.Dense(units=dims, activation=tf.nn.relu) def call(self, x_enc, enc_mask, x_dec, dec_mask, context, training=False, **kwargs): - with tf.name_scope("RecurrentDecoder"): - enc_mask = tf.expand_dims(enc_mask, axis=-1) - dec_mask = tf.expand_dims(dec_mask, axis=-1) - - initial_state = kwargs.get("initial_state") - x = self.drop(x_dec * dec_mask, training=training) - if initial_state is not None: - x = self.bi_lstm(x * dec_mask, initial_states=initial_state[0], training=training) - else: - x = self.bi_lstm(x * dec_mask, training=training) - x = self.h_drop(self.projection(x, projection_vector=context), training=training) - x = self.dense(x * dec_mask) - return x + enc_mask = tf.expand_dims(enc_mask, axis=-1) + dec_mask = tf.expand_dims(dec_mask, axis=-1) + + initial_state = kwargs.get("initial_state") + x = self.drop(x_dec * dec_mask, training=training) + if initial_state is not None: + x = self.bi_lstm(x * dec_mask, initial_states=initial_state[0], training=training) + else: + x = self.bi_lstm(x * dec_mask, training=training) + x = self.h_drop(self.projection(x, projection_vector=context), training=training) + x = self.dense(x * dec_mask) + return x diff --git a/text2vec/models/transformer.py b/text2vec/models/transformer.py index 1643240..2515825 100644 --- a/text2vec/models/transformer.py +++ b/text2vec/models/transformer.py @@ -1,4 +1,5 @@ import tensorflow as tf +from tensorflow.keras import layers from .components.attention import BahdanauAttention from .components.attention import MultiHeadAttention @@ -8,7 +9,7 @@ from .components.utils import TensorProjection -class TransformerEncoder(tf.keras.layers.Layer): +class TransformerEncoder(layers.Layer): """Attention based encoding pipeline. Parameters @@ -52,8 +53,8 @@ def __init__(self, max_sequence_len, layers=8, n_stacks=1, embedding_size=50, dims = embedding_size keep_prob = hidden_keep_prob - self.drop = tf.keras.layers.Dropout(1 - input_keep_prob, name="InputDropout") - self.h_drop = tf.keras.layers.Dropout(1 - hidden_keep_prob, name="HiddenStateDropout") + self.drop = layers.Dropout(1 - input_keep_prob, name="InputDropout") + self.h_drop = layers.Dropout(1 - hidden_keep_prob, name="HiddenStateDropout") self.layer_norm = LayerNorm() self.positional_encode = PositionalEncoder(emb_dims=dims, max_sequence_len=max_sequence_len) @@ -62,22 +63,20 @@ def __init__(self, max_sequence_len, layers=8, n_stacks=1, embedding_size=50, self.attention = BahdanauAttention(size=dims) def call(self, x, mask, training=False): - with tf.name_scope("TransformerEncoder"): - x = self.positional_encode(x, mask) - x = self.drop(x, training=training) - # mask = tf.expand_dims(mask, axis=-1) + x = self.positional_encode(x, mask) + x = self.drop(x, training=training) - for mha, ffn in zip(self.MHA, self.FFN): - x = self.h_drop(mha([x] * 3, training=training), training=training) + x - x = self.layer_norm(x) - x = self.h_drop(ffn(x), training=training) + x - x = self.layer_norm(x) + for mha, ffn in zip(self.MHA, self.FFN): + x = self.h_drop(mha([x] * 3, training=training), training=training) + x + x = self.layer_norm(x) + x = self.h_drop(ffn(x), training=training) + x + x = self.layer_norm(x) - context = self.attention(x) - return x, context + context = self.attention(x) + return x, context -class TransformerDecoder(tf.keras.layers.Layer): +class TransformerDecoder(layers.Layer): """Attention based decoding pipeline. Parameters @@ -102,8 +101,8 @@ def __init__(self, max_sequence_len, layers=8, n_stacks=1, embedding_size=50, dims = embedding_size keep_prob = hidden_keep_prob - self.drop = tf.keras.layers.Dropout(1 - input_keep_prob, name="InputDropout") - self.h_drop = tf.keras.layers.Dropout(1 - hidden_keep_prob, name="HiddenStateDropout") + self.drop = layers.Dropout(1 - input_keep_prob, name="InputDropout") + self.h_drop = layers.Dropout(1 - hidden_keep_prob, name="HiddenStateDropout") self.layer_norm = LayerNorm() self.projection = TensorProjection() @@ -112,25 +111,22 @@ def __init__(self, max_sequence_len, layers=8, n_stacks=1, embedding_size=50, self.FFN = [PositionWiseFFN(emb_dims=dims) for _ in range(n_stacks)] def call(self, x_enc, enc_mask, x_dec, dec_mask, context, attention, training=False, **kwargs): - with tf.name_scope("TransformerDecoder"): - x_dec = self.positional_encode(x_dec, dec_mask) - x_dec = self.drop(x_dec, training=training) - # enc_mask = tf.expand_dims(enc_mask, axis=-1) - # dec_mask = tf.expand_dims(dec_mask, axis=-1) - - for mha, ffn in zip(self.MHA, self.FFN): - x_dec = self.h_drop(mha( - [x_dec] * 3, - mask_future=True, - training=training - ), training=training) + x_dec - x_dec = self.layer_norm(x_dec) - - cross_context = attention(encoded=x_enc, decoded=x_dec) - x_dec = self.h_drop(self.projection(x_dec, projection_vector=cross_context), training=training) + x_dec - - x_dec = self.layer_norm(x_dec) - x_dec = self.h_drop(ffn(x_dec), training=training) + x_dec - x_dec = self.layer_norm(x_dec) - x_dec = self.h_drop(self.projection(x_dec, projection_vector=context), training=training) + x_dec - return x_dec + x_dec = self.positional_encode(x_dec, dec_mask) + x_dec = self.drop(x_dec, training=training) + + for mha, ffn in zip(self.MHA, self.FFN): + x_dec = self.h_drop(mha( + [x_dec] * 3, + mask_future=True, + training=training + ), training=training) + x_dec + x_dec = self.layer_norm(x_dec) + + cross_context = attention(encoded=x_enc, decoded=x_dec) + x_dec = self.h_drop(self.projection(x_dec, projection_vector=cross_context), training=training) + x_dec + + x_dec = self.layer_norm(x_dec) + x_dec = self.h_drop(ffn(x_dec), training=training) + x_dec + x_dec = self.layer_norm(x_dec) + x_dec = self.h_drop(self.projection(x_dec, projection_vector=context), training=training) + x_dec + return x_dec From 44c9123b3bff04530983e13cee47c6d0c829145f Mon Sep 17 00:00:00 2001 From: Dave Date: Wed, 6 Jul 2022 09:36:39 -0400 Subject: [PATCH 02/10] typing refactor --- text2vec/autoencoders.py | 47 ++++++++++++++++++---------------- text2vec/models/transformer.py | 13 +++++----- 2 files changed, 31 insertions(+), 29 deletions(-) diff --git a/text2vec/autoencoders.py b/text2vec/autoencoders.py index 8997ecc..497d1a9 100644 --- a/text2vec/autoencoders.py +++ b/text2vec/autoencoders.py @@ -1,7 +1,8 @@ # pylint: disable=too-many-ancestors -from typing import Dict +from typing import Dict, Optional import tensorflow as tf +from tensorflow.keras import layers, Model from text2vec.models.components.feeder import Tokenizer from text2vec.models.components.text_inputs import TokenEmbed @@ -12,7 +13,7 @@ from text2vec.models.sequential import RecurrentDecoder -class TransformerAutoEncoder(tf.keras.Model): +class TransformerAutoEncoder(Model): """Wrapper model class to combine the transformer based encoder-decoder training pipeline. Parameters @@ -27,7 +28,7 @@ class TransformerAutoEncoder(tf.keras.Model): Size of the vocabulary. Set this if pre-computing token IDs to pass to the model, by default None unknown_token : str, optional The placeholder value for OOV terms, by default '' - sep : int, optional + sep : str, optional Token separator by default ' ' input_keep_prob : float, optional Value between 0 and 1.0 which determines `1 - dropout_rate`, by default 1.0 @@ -54,7 +55,8 @@ class TransformerAutoEncoder(tf.keras.Model): """ def __init__(self, max_sequence_len: int, embedding_size: int, - token_hash: dict = None, vocab_size: int = None, unknown_token: str = '', sep: int = ' ', + token_hash: Optional[dict] = None, vocab_size: Optional[int] = None, + unknown_token: str = '', sep: str = ' ', input_keep_prob: float = 1.0, hidden_keep_prob: float = 1.0): super().__init__() @@ -77,23 +79,23 @@ def __init__(self, max_sequence_len: int, embedding_size: int, unknown_token=unknown_token ) else: - self.tokenizer = tf.keras.layers.Lambda(lambda x: x) # this is only for consistency, identity map + self.tokenizer = layers.Lambda(lambda x: x) # this is only for consistency, identity map self.embed_layer = Embed( vocab_size=vocab_size, embedding_size=embedding_size, max_sequence_len=max_sequence_len ) - self.encode_layer = TransformerEncoder(n_stacks=1, layers=8, **params) - self.decode_layer = TransformerDecoder(n_stacks=1, layers=8, **params) + self.encode_layer = TransformerEncoder(n_stacks=1, num_layers=8, **params) + self.decode_layer = TransformerDecoder(n_stacks=1, num_layers=8, **params) - def call(self, tokens, **kwargs): + def call(self, tokens, training: bool = False): # pylint: disable=missing-function-docstring tokens = self.tokenizer(tokens) - x_enc, enc_mask, _ = self.embed_layer(tokens, **kwargs) - x_enc, context = self.encode_layer(x_enc, mask=enc_mask, training=kwargs.get("training", False)) + x_enc, enc_mask, _ = self.embed_layer(tokens, training=training) + x_enc, context = self.encode_layer(x_enc, mask=enc_mask, training=training) return x_enc, context, enc_mask - def train_step(self, data): + def train_step(self, data): # pylint: disable=missing-function-docstring encoding_tok, decoding_tok = data decoding_tok = self.tokenizer(decoding_tok) @@ -134,7 +136,7 @@ def train_step(self, data): return {"loss": loss, 'learning_rate': self.optimizer.learning_rate(self.optimizer.iterations)} return {"loss": loss, 'learning_rate': self.optimizer.learning_rate} - def test_step(self, data): + def test_step(self, data): # pylint: disable=missing-function-docstring encoding_tok, decoding_tok = data decoding_tok = self.tokenizer(decoding_tok) @@ -208,7 +210,7 @@ def token_embed(self, sentences) -> Dict[str, tf.Tensor]: } -class LstmAutoEncoder(tf.keras.Model): +class LstmAutoEncoder(Model): """Wrapper model class to combine the LSTM based encoder-decoder training pipeline. Parameters @@ -225,7 +227,7 @@ class LstmAutoEncoder(tf.keras.Model): Size of the vocabulary. Set this if pre-computing token IDs to pass to the model, by default None unknown_token : str, optional The placeholder value for OOV terms, by default '' - sep : int, optional + sep : str, optional Token separator by default ' ' input_keep_prob : float, optional Value between 0 and 1.0 which determines `1 - dropout_rate`, by default 1.0 @@ -252,7 +254,8 @@ class LstmAutoEncoder(tf.keras.Model): """ def __init__(self, max_sequence_len: int, embedding_size: int, num_hidden: int = 64, - token_hash: dict = None, vocab_size: int = None, unknown_token: str = '', sep: int = ' ', + token_hash: Optional[dict] = None, vocab_size: Optional[int] = None, + unknown_token: str = '', sep: str = ' ', input_keep_prob: float = 1.0, hidden_keep_prob: float = 1.0): super().__init__() @@ -275,7 +278,7 @@ def __init__(self, max_sequence_len: int, embedding_size: int, num_hidden: int = unknown_token=unknown_token ) else: - self.tokenizer = tf.keras.layers.Lambda(lambda x: x) # this is only for consistency, identity map + self.tokenizer = layers.Lambda(lambda x: x) # this is only for consistency, identity map self.embed_layer = Embed( vocab_size=vocab_size, embedding_size=embedding_size, @@ -285,13 +288,13 @@ def __init__(self, max_sequence_len: int, embedding_size: int, num_hidden: int = self.encode_layer = RecurrentEncoder(num_hidden=num_hidden, **params) self.decode_layer = RecurrentDecoder(num_hidden=num_hidden, **params) - def call(self, tokens, **kwargs): + def call(self, tokens, training: bool = False): # pylint: disable=missing-function-docstring tokens = self.tokenizer(tokens) - x_enc, enc_mask, _ = self.embed_layer(tokens, **kwargs) - x_enc, context, *states = self.encode_layer(x_enc, mask=enc_mask, training=kwargs.get("training", False)) + x_enc, enc_mask, _ = self.embed_layer(tokens, training=training) + x_enc, context, *states = self.encode_layer(x_enc, mask=enc_mask, training=training) return x_enc, context, enc_mask, states - def train_step(self, data): + def train_step(self, data): # pylint: disable=missing-function-docstring encoding_tok, decoding_tok = data decoding_tok = self.tokenizer(decoding_tok) @@ -332,7 +335,7 @@ def train_step(self, data): return {"loss": loss, 'learning_rate': self.optimizer.learning_rate(self.optimizer.iterations)} return {"loss": loss, 'learning_rate': self.optimizer.learning_rate} - def test_step(self, data): + def test_step(self, data): # pylint: disable=missing-function-docstring encoding_tok, decoding_tok = data decoding_tok = self.tokenizer(decoding_tok) @@ -383,7 +386,7 @@ def embed(self, sentences) -> Dict[str, tf.Tensor]: and (batch_size, max_sequence_len, embedding_size) respectively. """ - sequences, attention, *args = self(sentences, training=False) + sequences, attention, *_ = self(sentences, training=False) return {"sequences": sequences, "attention": attention} @tf.function(input_signature=[tf.TensorSpec(shape=[None], dtype=tf.string)]) diff --git a/text2vec/models/transformer.py b/text2vec/models/transformer.py index 2515825..9336597 100644 --- a/text2vec/models/transformer.py +++ b/text2vec/models/transformer.py @@ -1,4 +1,3 @@ -import tensorflow as tf from tensorflow.keras import layers from .components.attention import BahdanauAttention @@ -16,7 +15,7 @@ class TransformerEncoder(layers.Layer): ---------- max_sequence_len : int Longest sequence seen at training time. - layers : int, optional + num_layers : int, optional Number of layers in the multi-head-attention layer, by default 8 n_stacks : int, optional Number of encoding blocks to chain, by default 1 @@ -47,7 +46,7 @@ class TransformerEncoder(layers.Layer): ``` """ - def __init__(self, max_sequence_len, layers=8, n_stacks=1, embedding_size=50, + def __init__(self, max_sequence_len, num_layers=8, n_stacks=1, embedding_size=50, input_keep_prob=1.0, hidden_keep_prob=1.0): super().__init__() dims = embedding_size @@ -58,7 +57,7 @@ def __init__(self, max_sequence_len, layers=8, n_stacks=1, embedding_size=50, self.layer_norm = LayerNorm() self.positional_encode = PositionalEncoder(emb_dims=dims, max_sequence_len=max_sequence_len) - self.MHA = [MultiHeadAttention(emb_dims=dims, layers=layers, keep_prob=keep_prob) for _ in range(n_stacks)] + self.MHA = [MultiHeadAttention(emb_dims=dims, layers=num_layers, keep_prob=keep_prob) for _ in range(n_stacks)] self.FFN = [PositionWiseFFN(emb_dims=dims) for _ in range(n_stacks)] self.attention = BahdanauAttention(size=dims) @@ -83,7 +82,7 @@ class TransformerDecoder(layers.Layer): ---------- max_sequence_len : int Longest sequence seen at training time. - layers : int, optional + num_layers : int, optional Number of layers in the multi-head-attention layer, by default 8 n_stacks : int, optional Number of encoding blocks to chain, by default 1 @@ -95,7 +94,7 @@ class TransformerDecoder(layers.Layer): Value between 0 and 1.0 which determines `1 - dropout_rate`, by default 1.0. """ - def __init__(self, max_sequence_len, layers=8, n_stacks=1, embedding_size=50, + def __init__(self, max_sequence_len, num_layers=8, n_stacks=1, embedding_size=50, input_keep_prob=1.0, hidden_keep_prob=1.0): super().__init__() dims = embedding_size @@ -107,7 +106,7 @@ def __init__(self, max_sequence_len, layers=8, n_stacks=1, embedding_size=50, self.projection = TensorProjection() self.positional_encode = PositionalEncoder(emb_dims=dims, max_sequence_len=max_sequence_len) - self.MHA = [MultiHeadAttention(emb_dims=dims, layers=layers, keep_prob=keep_prob) for _ in range(n_stacks)] + self.MHA = [MultiHeadAttention(emb_dims=dims, layers=num_layers, keep_prob=keep_prob) for _ in range(n_stacks)] self.FFN = [PositionWiseFFN(emb_dims=dims) for _ in range(n_stacks)] def call(self, x_enc, enc_mask, x_dec, dec_mask, context, attention, training=False, **kwargs): From 150b65e67478be4fe6927a22800530347eaa17cc Mon Sep 17 00:00:00 2001 From: Dave Date: Wed, 6 Jul 2022 10:33:46 -0400 Subject: [PATCH 03/10] change keep rates to drop rates --- examples/trainers/wiki_lstm.py | 4 +- examples/trainers/wiki_transformer.py | 4 +- tests/lstm_auto_enc_test_fit.py | 4 +- tests/transformer_auto_enc_test_fit.py | 4 +- text2vec/autoencoders.py | 29 ++++---- text2vec/models/components/attention.py | 88 +++++++++++++------------ text2vec/models/sequential.py | 28 ++++---- text2vec/models/transformer.py | 50 +++++++------- 8 files changed, 111 insertions(+), 100 deletions(-) diff --git a/examples/trainers/wiki_lstm.py b/examples/trainers/wiki_lstm.py index 60f9a30..faffaf2 100644 --- a/examples/trainers/wiki_lstm.py +++ b/examples/trainers/wiki_lstm.py @@ -99,8 +99,8 @@ def token_mapper(text: Union[str, List[str]]): max_sequence_len=512, embedding_size=128, token_hash=tokenizer.get_vocab(), - input_keep_prob=0.7, - hidden_keep_prob=0.5 + input_drop_rate=0.3, + hidden_drop_rate=0.5 ) model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=RampUpDecaySchedule(embedding_size=128))) checkpoint = tf.train.Checkpoint(Classifier=model, optimizer=model.optimizer) diff --git a/examples/trainers/wiki_transformer.py b/examples/trainers/wiki_transformer.py index f0dde3e..06dbeb9 100644 --- a/examples/trainers/wiki_transformer.py +++ b/examples/trainers/wiki_transformer.py @@ -99,8 +99,8 @@ def token_mapper(text: Union[str, List[str]]): max_sequence_len=512, embedding_size=128, token_hash=tokenizer.get_vocab(), - input_keep_prob=0.7, - hidden_keep_prob=0.5 + input_drop_rate=0.3, + hidden_drop_rate=0.5 ) model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=RampUpDecaySchedule(embedding_size=128))) checkpoint = tf.train.Checkpoint(Classifier=model, optimizer=model.optimizer) diff --git a/tests/lstm_auto_enc_test_fit.py b/tests/lstm_auto_enc_test_fit.py index c56da83..4950472 100644 --- a/tests/lstm_auto_enc_test_fit.py +++ b/tests/lstm_auto_enc_test_fit.py @@ -92,8 +92,8 @@ def token_mapper(text: Union[str, List[str]]): max_sequence_len=256, embedding_size=256, token_hash=tokenizer.get_vocab(), - input_keep_prob=0.7, - hidden_keep_prob=0.5 + input_drop_rate=0.3, + hidden_drop_rate=0.5 ) model.compile(optimizer=tf.keras.optimizers.Adam(0.01), run_eagerly=True) model.fit(x=data.prefetch(10).batch(16), epochs=1) diff --git a/tests/transformer_auto_enc_test_fit.py b/tests/transformer_auto_enc_test_fit.py index a4c8265..33953af 100644 --- a/tests/transformer_auto_enc_test_fit.py +++ b/tests/transformer_auto_enc_test_fit.py @@ -92,8 +92,8 @@ def token_mapper(text: Union[str, List[str]]): max_sequence_len=512, embedding_size=128, token_hash=tokenizer.get_vocab(), - input_keep_prob=0.7, - hidden_keep_prob=0.5 + input_drop_rate=0.3, + hidden_drop_rate=0.5 ) model.compile(optimizer=tf.keras.optimizers.Adam(), run_eagerly=True) model.fit(x=data.prefetch(10).batch(16), epochs=1) diff --git a/text2vec/autoencoders.py b/text2vec/autoencoders.py index 497d1a9..10ef352 100644 --- a/text2vec/autoencoders.py +++ b/text2vec/autoencoders.py @@ -30,10 +30,10 @@ class TransformerAutoEncoder(Model): The placeholder value for OOV terms, by default '' sep : str, optional Token separator by default ' ' - input_keep_prob : float, optional - Value between 0 and 1.0 which determines `1 - dropout_rate`, by default 1.0 - hidden_keep_prob : float, optional - Hidden states dropout. Value between 0 and 1.0 which determines `1 - dropout_rate`, by default 1.0 + input_drop_rate : float, optional + Value between 0 and 1.0, by default 0. + hidden_drop_rate : float, optional + Value between 0 and 1.0, by default 0. Raises ------ @@ -57,7 +57,7 @@ class TransformerAutoEncoder(Model): def __init__(self, max_sequence_len: int, embedding_size: int, token_hash: Optional[dict] = None, vocab_size: Optional[int] = None, unknown_token: str = '', sep: str = ' ', - input_keep_prob: float = 1.0, hidden_keep_prob: float = 1.0): + input_drop_rate: float = 0, hidden_drop_rate: float = 0): super().__init__() if token_hash is None and vocab_size is None: @@ -66,8 +66,8 @@ def __init__(self, max_sequence_len: int, embedding_size: int, params = dict( max_sequence_len=max_sequence_len, embedding_size=embedding_size, - input_keep_prob=input_keep_prob, - hidden_keep_prob=hidden_keep_prob + input_drop_rate=input_drop_rate, + hidden_drop_rate=hidden_drop_rate ) if token_hash is not None: @@ -112,7 +112,6 @@ def train_step(self, data): # pylint: disable=missing-function-docstring decoding_tok, dec_mask, _ = self.embed_layer(decoding_tok[:, :-1]) # skip decoding_tok = self.decode_layer( x_enc=x_enc, - enc_mask=enc_mask, x_dec=decoding_tok, dec_mask=dec_mask, context=context, @@ -229,10 +228,10 @@ class LstmAutoEncoder(Model): The placeholder value for OOV terms, by default '' sep : str, optional Token separator by default ' ' - input_keep_prob : float, optional - Value between 0 and 1.0 which determines `1 - dropout_rate`, by default 1.0 - hidden_keep_prob : float, optional - Hidden states dropout. Value between 0 and 1.0 which determines `1 - dropout_rate`, by default 1.0 + input_drop_rate : float, optional + Value between 0 and 1.0, by default 0. + hidden_drop_rate : float, optional + Value between 0 and 1.0, by default 0. Raises ------ @@ -256,7 +255,7 @@ class LstmAutoEncoder(Model): def __init__(self, max_sequence_len: int, embedding_size: int, num_hidden: int = 64, token_hash: Optional[dict] = None, vocab_size: Optional[int] = None, unknown_token: str = '', sep: str = ' ', - input_keep_prob: float = 1.0, hidden_keep_prob: float = 1.0): + input_drop_rate: float = 0., hidden_drop_rate: float = 0.): super().__init__() if token_hash is None and vocab_size is None: @@ -265,8 +264,8 @@ def __init__(self, max_sequence_len: int, embedding_size: int, num_hidden: int = params = dict( max_sequence_len=max_sequence_len, embedding_size=embedding_size, - input_keep_prob=input_keep_prob, - hidden_keep_prob=hidden_keep_prob + input_drop_rate=input_drop_rate, + hidden_drop_rate=hidden_drop_rate ) if token_hash is not None: diff --git a/text2vec/models/components/attention.py b/text2vec/models/components/attention.py index bfdf224..7a4d2f8 100644 --- a/text2vec/models/components/attention.py +++ b/text2vec/models/components/attention.py @@ -38,19 +38,20 @@ class ScaledDotAttention(layers.Layer): def __init__(self): super().__init__(name="ScaledDotAttention") + self.scale = tf.constant(1e9, dtype=tf.float32) - def call(self, query, key, value, mask_future=False): - with tf.name_scope("ScaledDotAttention"): - numerator = tf.einsum('ijk,ilk->ijl', query, key) - denominator = tf.sqrt(tf.cast(tf.shape(key)[-1], tf.float32)) + # pylint: disable=missing-function-docstring + def call(self, query, key, value, mask_future: bool = False): + numerator = tf.einsum('ijk,ilk->ijl', query, key) + denominator = tf.sqrt(tf.cast(tf.shape(key)[-1], tf.float32)) - if mask_future: - upper = (1 + 1e9) * tf.linalg.band_part(tf.ones_like(numerator), num_lower=0, num_upper=-1) - mask = 1 - upper - numerator *= mask + if mask_future: + upper = (1 + self.scale) * tf.linalg.band_part(tf.ones_like(numerator), num_lower=0, num_upper=-1) + mask = 1 - upper + numerator *= mask - x = tf.nn.softmax(numerator / denominator) - return tf.einsum('ijk,ikl->ijl', x, value) + x = tf.nn.softmax(numerator / denominator) + return tf.einsum('ijk,ikl->ijl', x, value) class BahdanauAttention(layers.Layer): @@ -65,6 +66,8 @@ class BahdanauAttention(layers.Layer): ---------- size : int The dimensionality of the hidden attention weights. This is the same as the word-embedding dimensionality. + drop_rate : float, optional + Value between 0 and 1.0, performs dropout on the attention weights, by default 0. Examples -------- @@ -85,13 +88,14 @@ class BahdanauAttention(layers.Layer): ``` """ - def __init__(self, size: int, dropout_rate: float = 0.): + def __init__(self, size: int, drop_rate: float = 0.): super().__init__(name="BahdanauAttention") self.hidden = layers.Dense(units=size, activation="tanh") self.U = tf.Variable(initializers.GlorotUniform()(shape=[size]), name="U", dtype=tf.float32, trainable=True) - self.dropout = layers.Dropout(dropout_rate) + self.dropout = layers.Dropout(drop_rate) + # pylint: disable=missing-function-docstring def call(self, encoded: tf.Tensor, decoded: Optional[tf.Tensor] = None, training: bool = False) -> tf.Tensor: if decoded is None: score = tf.math.reduce_sum(self.U * self.hidden(encoded), axis=-1) @@ -103,6 +107,7 @@ def call(self, encoded: tf.Tensor, decoded: Optional[tf.Tensor] = None, training score = tf.einsum("ijm,mn,ikn->ijk", encoded, self.hidden.kernel, decoded) alphas = tf.nn.softmax(score, axis=1) + alphas = self.dropout(alphas, training=training) alphas = tf.math.reduce_sum(tf.matmul(alphas, encoded, transpose_a=True), axis=-1) return tf.einsum('ilk,il->ik', decoded, alphas) @@ -120,10 +125,10 @@ class SingleHeadAttention(layers.Layer): ---------- emb_dims : int The word-embedding dimensionality. This value determines the dimensionalities of the hidden weights. - layers : int, optional + num_layers : int, optional The number of parallel single-head-attention mechanisms, by default 8. - keep_prob : float, optional - Value between 0 and 1.0 which determines `1 - dropout_rate`, by default 1.0. + drop_rate : float, optional + Value between 0 and 1.0, by default 0. Examples -------- @@ -143,32 +148,32 @@ class SingleHeadAttention(layers.Layer): ``` """ - def __init__(self, emb_dims, layers=8, keep_prob=1.0): + def __init__(self, emb_dims, num_layers: int = 8, drop_rate: float = 0.): super().__init__(name="SingleHeadAttention") - assert isinstance(layers, int) and layers > 0 + assert isinstance(num_layers, int) and num_layers > 0 dims = emb_dims - key_dims = emb_dims // layers + key_dims = emb_dims // num_layers initializer = tf.keras.initializers.GlorotUniform() self.WQ = tf.Variable(initializer(shape=(dims, key_dims)), name="WQ", dtype=tf.float32, trainable=True) self.WK = tf.Variable(initializer(shape=(dims, key_dims)), name="WK", dtype=tf.float32, trainable=True) self.WV = tf.Variable(initializer(shape=(dims, key_dims)), name="WV", dtype=tf.float32, trainable=True) - self.dropout = layers.Dropout(1 - keep_prob) + self.dropout = layers.Dropout(drop_rate) self.dot_attention = ScaledDotAttention() + # pylint: disable=missing-function-docstring def call(self, inputs, mask_future: bool = False, training: bool = False): - with tf.name_scope("SingleHeadAttention"): - queries, keys, values = inputs + queries, keys, values = inputs - queries = self.dropout(queries, training=training) - keys = self.dropout(keys, training=training) - values = self.dropout(values, training=training) + queries = self.dropout(queries, training=training) + keys = self.dropout(keys, training=training) + values = self.dropout(values, training=training) - head_queries = tf.tensordot(queries, self.WQ, axes=[-1, 0]) - head_keys = tf.tensordot(keys, self.WK, axes=[-1, 0]) - head_values = tf.tensordot(values, self.WV, axes=[-1, 0]) - return self.dot_attention(query=head_queries, key=head_keys, value=head_values, mask_future=mask_future) + head_queries = tf.tensordot(queries, self.WQ, axes=[-1, 0]) + head_keys = tf.tensordot(keys, self.WK, axes=[-1, 0]) + head_values = tf.tensordot(values, self.WV, axes=[-1, 0]) + return self.dot_attention(query=head_queries, key=head_keys, value=head_values, mask_future=mask_future) class MultiHeadAttention(layers.Layer): @@ -184,10 +189,10 @@ class MultiHeadAttention(layers.Layer): ---------- emb_dims : int The word-embedding dimensionality. This value determines the dimensionalities of the hidden weights. - layers : int, optional + num_layers : int, optional The number of parallel single-head-attention mechanisms, by default 8. - keep_prob : float, optional - Value between 0 and 1.0 which determines `1 - dropout_rate`, by default 1.0. + drop_rate : float, optional + Value between 0 and 1.0, by default 0. Examples -------- @@ -207,17 +212,16 @@ class MultiHeadAttention(layers.Layer): ``` """ - def __init__(self, emb_dims: int, layers: int = 8, keep_prob: float = 1.0): + def __init__(self, emb_dims: int, num_layers: int = 8, drop_rate: float = 0.): super().__init__(name="MultiHeadAttention") - self.layer_heads = [] - for i in range(layers): - with tf.name_scope(f"head-{i}"): - self.layer_heads.append(SingleHeadAttention(emb_dims=emb_dims, layers=layers, keep_prob=keep_prob)) - - self.dense =layers.Dense(units=emb_dims, use_bias=False) + self.layer_heads = [ + SingleHeadAttention(emb_dims=emb_dims, num_layers=num_layers, drop_rate=drop_rate) + for _ in range(num_layers) + ] + self.dense = layers.Dense(units=emb_dims, use_bias=False) + # pylint: disable=missing-function-docstring def call(self, inputs, mask_future=False, training=False): - with tf.name_scope("MultiHeadAttention"): - heads = [layer(inputs, mask_future=mask_future, training=training) for layer in self.layer_heads] - total_head = tf.concat(heads, -1) - return self.dense(total_head) + heads = [layer(inputs, mask_future=mask_future, training=training) for layer in self.layer_heads] + total_head = tf.concat(heads, -1) + return self.dense(total_head) diff --git a/text2vec/models/sequential.py b/text2vec/models/sequential.py index 46a5869..1f533ea 100644 --- a/text2vec/models/sequential.py +++ b/text2vec/models/sequential.py @@ -17,8 +17,8 @@ class RecurrentEncoder(layers.Layer): Dimensionality of hidden LSTM layer weights. num_layers : int, optional Number of hidden LSTM layers, by default 2 - input_keep_prob : float, optional - Value between 0 and 1.0 which determines `1 - dropout_rate`, by default 1.0. + input_drop_rate : float, optional + Value between 0 and 1.0, by default 0. Examples -------- @@ -40,15 +40,16 @@ class RecurrentEncoder(layers.Layer): ``` """ - def __init__(self, max_sequence_len, num_hidden, num_layers=2, input_keep_prob=1.0, **kwargs): + def __init__(self, max_sequence_len, num_hidden, num_layers=2, input_drop_rate: float = 0., **kwargs): super().__init__() self.max_sequence_length = max_sequence_len - self.drop = layers.Dropout(1 - input_keep_prob, name="InputDropout") + self.drop = layers.Dropout(input_drop_rate) self.bi_lstm = BidirectionalLSTM(num_layers=num_layers, num_hidden=num_hidden, return_states=True) self.attention = BahdanauAttention(size=2 * num_hidden) - def call(self, x, mask, training=False, **kwargs): + # pylint: disable=missing-function-docstring + def call(self, x, mask, training: bool = False): with tf.name_scope("RecurrentEncoder"): mask = tf.expand_dims(mask, axis=-1) x = self.drop(x, training=training) @@ -73,26 +74,27 @@ class RecurrentDecoder(layers.Layer): Dimensionality of the word-embeddings, by default 50. num_layers : int, optional Number of hidden LSTM layers, by default 2 - input_keep_prob : float, optional - Value between 0 and 1.0 which determines `1 - dropout_rate`, by default 1.0. - hidden_keep_prob : float, optional - Hidden states dropout. Value between 0 and 1.0 which determines `1 - dropout_rate`, by default 1.0. + input_drop_rate : float, optional + Value between 0 and 1.0, by default 0. + hidden_drop_rate : float, optional + Value between 0 and 1.0, by default 0. """ def __init__(self, max_sequence_len, num_hidden, embedding_size=50, num_layers=2, - input_keep_prob=1.0, hidden_keep_prob=1.0): + input_drop_rate: float = 0., hidden_drop_rate: float = 0.): super().__init__() self.max_sequence_length = max_sequence_len dims = embedding_size - self.drop = layers.Dropout(1 - input_keep_prob, name="InputDropout") - self.h_drop = layers.Dropout(1 - hidden_keep_prob, name="HiddenStateDropout") + self.drop = layers.Dropout(input_drop_rate) + self.h_drop = layers.Dropout(hidden_drop_rate) self.projection = TensorProjection() self.bi_lstm = BidirectionalLSTM(num_layers=num_layers, num_hidden=num_hidden, return_states=False) self.dense = layers.Dense(units=dims, activation=tf.nn.relu) - def call(self, x_enc, enc_mask, x_dec, dec_mask, context, training=False, **kwargs): + # pylint: disable=missing-function-docstring + def call(self, x_enc, enc_mask, x_dec, dec_mask, context, training: bool = False): enc_mask = tf.expand_dims(enc_mask, axis=-1) dec_mask = tf.expand_dims(dec_mask, axis=-1) diff --git a/text2vec/models/transformer.py b/text2vec/models/transformer.py index 9336597..5546c66 100644 --- a/text2vec/models/transformer.py +++ b/text2vec/models/transformer.py @@ -21,10 +21,10 @@ class TransformerEncoder(layers.Layer): Number of encoding blocks to chain, by default 1 embedding_size : int, optional Dimensionality of the word-embeddings, by default 50. - input_keep_prob : float, optional - Value between 0 and 1.0 which determines `1 - dropout_rate`, by default 1.0. - hidden_keep_prob : float, optional - Value between 0 and 1.0 which determines `1 - dropout_rate`, by default 1.0. + input_drop_rate : float, optional + Value between 0 and 1.0, by default 0. + hidden_drop_rate : float, optional + Value between 0 and 1.0, by default 0. Examples -------- @@ -47,21 +47,24 @@ class TransformerEncoder(layers.Layer): """ def __init__(self, max_sequence_len, num_layers=8, n_stacks=1, embedding_size=50, - input_keep_prob=1.0, hidden_keep_prob=1.0): + input_drop_rate: float = 0., hidden_drop_rate: float = 0.): super().__init__() dims = embedding_size - keep_prob = hidden_keep_prob - self.drop = layers.Dropout(1 - input_keep_prob, name="InputDropout") - self.h_drop = layers.Dropout(1 - hidden_keep_prob, name="HiddenStateDropout") + self.drop = layers.Dropout(input_drop_rate) + self.h_drop = layers.Dropout(hidden_drop_rate) self.layer_norm = LayerNorm() self.positional_encode = PositionalEncoder(emb_dims=dims, max_sequence_len=max_sequence_len) - self.MHA = [MultiHeadAttention(emb_dims=dims, layers=num_layers, keep_prob=keep_prob) for _ in range(n_stacks)] + self.MHA = [ + MultiHeadAttention(emb_dims=dims, num_layers=num_layers, drop_rate=input_drop_rate) + for _ in range(n_stacks) + ] self.FFN = [PositionWiseFFN(emb_dims=dims) for _ in range(n_stacks)] self.attention = BahdanauAttention(size=dims) - def call(self, x, mask, training=False): + # pylint: disable=missing-function-docstring + def call(self, x, mask, training: bool = False): x = self.positional_encode(x, mask) x = self.drop(x, training=training) @@ -71,7 +74,7 @@ def call(self, x, mask, training=False): x = self.h_drop(ffn(x), training=training) + x x = self.layer_norm(x) - context = self.attention(x) + x, context = self.attention(x) return x, context @@ -88,28 +91,31 @@ class TransformerDecoder(layers.Layer): Number of encoding blocks to chain, by default 1 embedding_size : int, optional Dimensionality of the word-embeddings, by default 50. - input_keep_prob : float, optional - Value between 0 and 1.0 which determines `1 - dropout_rate`, by default 1.0. - hidden_keep_prob : float, optional - Value between 0 and 1.0 which determines `1 - dropout_rate`, by default 1.0. + input_drop_rate : float, optional + Value between 0 and 1.0, by default 0. + hidden_drop_rate : float, optional + Value between 0 and 1.0, by default 0. """ def __init__(self, max_sequence_len, num_layers=8, n_stacks=1, embedding_size=50, - input_keep_prob=1.0, hidden_keep_prob=1.0): + input_drop_rate: float = 0., hidden_drop_rate: float = 0.): super().__init__() dims = embedding_size - keep_prob = hidden_keep_prob - self.drop = layers.Dropout(1 - input_keep_prob, name="InputDropout") - self.h_drop = layers.Dropout(1 - hidden_keep_prob, name="HiddenStateDropout") + self.drop = layers.Dropout(input_drop_rate) + self.h_drop = layers.Dropout(hidden_drop_rate) self.layer_norm = LayerNorm() self.projection = TensorProjection() self.positional_encode = PositionalEncoder(emb_dims=dims, max_sequence_len=max_sequence_len) - self.MHA = [MultiHeadAttention(emb_dims=dims, layers=num_layers, keep_prob=keep_prob) for _ in range(n_stacks)] + self.MHA = [ + MultiHeadAttention(emb_dims=dims, num_layers=num_layers, drop_rate=input_drop_rate) + for _ in range(n_stacks) + ] self.FFN = [PositionWiseFFN(emb_dims=dims) for _ in range(n_stacks)] - def call(self, x_enc, enc_mask, x_dec, dec_mask, context, attention, training=False, **kwargs): + # pylint: disable=missing-function-docstring + def call(self, x_enc, x_dec, dec_mask, context, attention: BahdanauAttention, training: bool = False): x_dec = self.positional_encode(x_dec, dec_mask) x_dec = self.drop(x_dec, training=training) @@ -121,7 +127,7 @@ def call(self, x_enc, enc_mask, x_dec, dec_mask, context, attention, training=Fa ), training=training) + x_dec x_dec = self.layer_norm(x_dec) - cross_context = attention(encoded=x_enc, decoded=x_dec) + x_dec, cross_context = attention(encoded=x_enc, decoded=x_dec) x_dec = self.h_drop(self.projection(x_dec, projection_vector=cross_context), training=training) + x_dec x_dec = self.layer_norm(x_dec) From 75a7c40333d85eba84b28276bd7bce97d252551b Mon Sep 17 00:00:00 2001 From: Dave Date: Wed, 6 Jul 2022 13:48:50 -0400 Subject: [PATCH 04/10] multi-news training example --- .gitignore | 1 + ...iki_transformer.py => news_transformer.py} | 87 ++++++++++--------- tests/transformer_auto_enc_test_fit.py | 16 ++-- text2vec/autoencoders.py | 2 +- text2vec/models/components/attention.py | 4 +- 5 files changed, 61 insertions(+), 49 deletions(-) rename examples/trainers/{wiki_transformer.py => news_transformer.py} (62%) diff --git a/.gitignore b/.gitignore index a498a45..55960d9 100644 --- a/.gitignore +++ b/.gitignore @@ -7,6 +7,7 @@ **/*.dev.yml wiki_t2v/ +multi_news_t2v/ # JavaScript configs and dependencies **/.eslintrc.json diff --git a/examples/trainers/wiki_transformer.py b/examples/trainers/news_transformer.py similarity index 62% rename from examples/trainers/wiki_transformer.py rename to examples/trainers/news_transformer.py index 06dbeb9..4081f65 100644 --- a/examples/trainers/wiki_transformer.py +++ b/examples/trainers/news_transformer.py @@ -13,13 +13,17 @@ import numpy as np import tensorflow as tf +from tensorflow.keras import optimizers, callbacks +from tensorflow.keras import backend as K from tensorboard.plugins import projector from text2vec.autoencoders import TransformerAutoEncoder from text2vec.optimizer_tools import RampUpDecaySchedule os.environ["TOKENIZERS_PARALLELISM"] = "true" -sent_tokenizer = PunktSentenceTokenizer().tokenize +root = os.path.dirname(os.path.abspath(__file__)) +EMBEDDING_SIZE = 128 +MAX_SEQUENCE_LENGTH = 512 def train_tokenizer() -> Tuple[tokenizers.Tokenizer, Generator, int]: @@ -41,11 +45,13 @@ def train_tokenizer() -> Tuple[tokenizers.Tokenizer, Generator, int]: special_tokens=[("[SEP]", 1), ("", 2), ("", 3)] ) - dataset = datasets.load_dataset("wikitext", "wikitext-103-raw-v1", split="test") + # dataset = datasets.load_dataset("wikitext", "wikitext-103-raw-v1", split="test") + dataset = datasets.load_dataset("multi_news", split="test") def batch_iterator(batch_size=1000): for i in range(0, len(dataset), batch_size): - yield dataset[i: i + batch_size]["text"] + for key in dataset.features: + yield dataset[i: i + batch_size][key] tokenizer.train_from_iterator( batch_iterator(), @@ -55,14 +61,32 @@ def batch_iterator(batch_size=1000): ) ) + tokenizer.enable_truncation(2 * MAX_SEQUENCE_LENGTH + 3) # 2 for the [SEP], , tokens + tokenizer.post_processor = tokenizers.processors.TemplateProcessing( + single="$A", + pair="$A:0 [SEP] $B:1 ", + special_tokens=[ + ("[SEP]", 1), + ("", 2), + ("", 3) + ] + ) + def generator(): for record in dataset: - if record['text'].strip() != '': - for sentence in sent_tokenizer(record['text']): - yield sentence - - data = tf.data.Dataset.from_generator(generator, output_signature=(tf.TensorSpec(shape=(None), dtype=tf.string))) - data = data.map(tf.strings.strip, num_parallel_calls=tf.data.experimental.AUTOTUNE) + if record["document"] and record["summary"]: + enc, dec = ' '.join(tokenizer.encode( + record["document"], + pair=record["summary"] + ).tokens).split(' [SEP] ', maxsplit=2) + + if enc.strip() != "" and dec != "": + yield enc, dec + + data = tf.data.Dataset.from_generator( + generator, + output_signature=(tf.TensorSpec(shape=(None), dtype=tf.string), tf.TensorSpec(shape=(None), dtype=tf.string)) + ) return tokenizer, data @@ -77,32 +101,16 @@ def main(save_path: str): for token, _ in sorted(tokenizer.get_vocab().items(), key=lambda s: s[-1]): tsv.write(f"{token}\n") - def encode(x): - def token_mapper(text: Union[str, List[str]]): - text = text.numpy() - - if isinstance(text, np.ndarray): - enc, dec = [], [] - for batch in tokenizer.encode_batch([(t.decode('utf8'), t.decode('utf8')) for t in text]): - enc_, dec_ = ' '.join(batch.tokens).split(' [SEP] ') - enc.append(enc_) - dec.append(dec_) - return (enc, dec) - - text = text.decode('utf8') - enc, dec = ' '.join(tokenizer.encode(text, pair=text).tokens).split(' [SEP] ') - return (enc, dec) - - return tf.py_function(token_mapper, inp=[x], Tout=[tf.string, tf.string]) - model = TransformerAutoEncoder( - max_sequence_len=512, - embedding_size=128, + max_sequence_len=MAX_SEQUENCE_LENGTH, + embedding_size=EMBEDDING_SIZE, token_hash=tokenizer.get_vocab(), input_drop_rate=0.3, hidden_drop_rate=0.5 ) - model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=RampUpDecaySchedule(embedding_size=128))) + + scheduler = RampUpDecaySchedule(EMBEDDING_SIZE, warmup_steps=4000) + model.compile(optimizer=optimizers.Adam(scheduler(0).numpy())) checkpoint = tf.train.Checkpoint(Classifier=model, optimizer=model.optimizer) checkpoint_manager = tf.train.CheckpointManager(checkpoint, save_path, max_to_keep=3) @@ -117,18 +125,19 @@ def token_mapper(text: Union[str, List[str]]): embeddings_config.metadata_path = f"{save_path}/metadata.tsv" projector.visualize_embeddings(logdir=save_path, config=config) - data = data.map(encode, num_parallel_calls=tf.data.experimental.AUTOTUNE) model.fit( x=data.prefetch(8).batch(64), callbacks=[ - tf.keras.callbacks.TensorBoard( - log_dir=save_path, - write_graph=True, - update_freq=100 - ), - tf.keras.callbacks.LambdaCallback(on_epoch_end=lambda epoch, logs: checkpoint_manager.save()) + callbacks.TensorBoard(log_dir=save_path, write_graph=True, update_freq=100), + callbacks.LambdaCallback( + on_epoch_end=lambda epoch, logs: checkpoint_manager.save(), + on_batch_end=lambda batch, logs: K.set_value( + model.optimizer.lr, + K.get_value(scheduler(model.optimizer.iterations)) + ) + ) ], - epochs=1 + epochs=10 ) model.save( @@ -142,4 +151,4 @@ def token_mapper(text: Union[str, List[str]]): if __name__ == '__main__': - main(save_path='./wiki_t2v') + main(save_path=f'{root}/../../multi_news_t2v') diff --git a/tests/transformer_auto_enc_test_fit.py b/tests/transformer_auto_enc_test_fit.py index 33953af..ed2c796 100644 --- a/tests/transformer_auto_enc_test_fit.py +++ b/tests/transformer_auto_enc_test_fit.py @@ -1,5 +1,5 @@ -import os from typing import List, Union +import os import datasets import tokenizers @@ -12,6 +12,7 @@ import numpy as np import tensorflow as tf +from tensorflow.keras import optimizers from text2vec.autoencoders import TransformerAutoEncoder root = os.path.dirname(os.path.abspath(__file__)) @@ -85,8 +86,8 @@ def token_mapper(text: Union[str, List[str]]): return tf.py_function(token_mapper, inp=[x], Tout=[tf.string, tf.string]) data = tf.data.Dataset.from_generator(data_gen, output_signature=(tf.TensorSpec(shape=(None), dtype=tf.string))) - data = data.map(tf.strings.strip, num_parallel_calls=tf.data.experimental.AUTOTUNE) - data = data.map(encode, num_parallel_calls=tf.data.experimental.AUTOTUNE) + data = data.map(tf.strings.strip, num_parallel_calls=tf.data.AUTOTUNE) + data = data.map(encode, num_parallel_calls=tf.data.AUTOTUNE) model = TransformerAutoEncoder( max_sequence_len=512, @@ -95,11 +96,12 @@ def token_mapper(text: Union[str, List[str]]): input_drop_rate=0.3, hidden_drop_rate=0.5 ) - model.compile(optimizer=tf.keras.optimizers.Adam(), run_eagerly=True) - model.fit(x=data.prefetch(10).batch(16), epochs=1) + model.compile(optimizer=optimizers.Adam(1e-4), run_eagerly=False) + model.fit(x=data.prefetch(10).batch(16), epochs=10) - model(['here is a sentence', 'try another one']) - model.predict(['here is a sentence', 'try another one']) + x = model.embed(['this is about physics', 'this is not about physics'])["attention"] + x = tf.linalg.l2_normalize(x, axis=-1) + print(x.numpy() @ x.numpy().T) return model diff --git a/text2vec/autoencoders.py b/text2vec/autoencoders.py index 10ef352..81c9295 100644 --- a/text2vec/autoencoders.py +++ b/text2vec/autoencoders.py @@ -101,7 +101,7 @@ def train_step(self, data): # pylint: disable=missing-function-docstring with tf.GradientTape() as tape: with tf.name_scope('Encoding'): - x_enc, context, enc_mask = self(encoding_tok, training=True) + x_enc, context, _ = self(encoding_tok, training=True) with tf.name_scope('Decoding'): targets = decoding_tok[:, 1:] # skip the token with the slice on axis=1 diff --git a/text2vec/models/components/attention.py b/text2vec/models/components/attention.py index 7a4d2f8..922821b 100644 --- a/text2vec/models/components/attention.py +++ b/text2vec/models/components/attention.py @@ -103,13 +103,13 @@ def call(self, encoded: tf.Tensor, decoded: Optional[tf.Tensor] = None, training alphas = self.dropout(alphas, training=training) x = tf.expand_dims(alphas, axis=-1) * encoded return x, tf.math.reduce_sum(x, axis=1) - # return tf.einsum('ilk,il->ik', encoded, alphas) score = tf.einsum("ijm,mn,ikn->ijk", encoded, self.hidden.kernel, decoded) alphas = tf.nn.softmax(score, axis=1) alphas = self.dropout(alphas, training=training) alphas = tf.math.reduce_sum(tf.matmul(alphas, encoded, transpose_a=True), axis=-1) - return tf.einsum('ilk,il->ik', decoded, alphas) + x = tf.expand_dims(alphas, axis=-1) * decoded + return x, tf.math.reduce_sum(x, axis=1) class SingleHeadAttention(layers.Layer): From d9052461073a537b78b080617327eb229dbae824 Mon Sep 17 00:00:00 2001 From: Dave Date: Thu, 7 Jul 2022 11:35:06 -0400 Subject: [PATCH 05/10] reduced ops --- examples/trainers/news_transformer.py | 20 +--- text2vec/autoencoders.py | 110 +++++++++------------- text2vec/models/components/attention.py | 12 +-- text2vec/models/components/text_inputs.py | 13 ++- text2vec/models/components/utils.py | 58 ++++++++++-- text2vec/models/transformer.py | 20 ++-- 6 files changed, 126 insertions(+), 107 deletions(-) diff --git a/examples/trainers/news_transformer.py b/examples/trainers/news_transformer.py index 4081f65..41408d6 100644 --- a/examples/trainers/news_transformer.py +++ b/examples/trainers/news_transformer.py @@ -1,4 +1,4 @@ -from typing import Generator, List, Tuple, Union +from typing import Generator, Tuple import os import datasets @@ -9,9 +9,7 @@ from tokenizers import pre_tokenizers from tokenizers import processors from tokenizers import trainers -from nltk.tokenize import PunktSentenceTokenizer -import numpy as np import tensorflow as tf from tensorflow.keras import optimizers, callbacks from tensorflow.keras import backend as K @@ -26,7 +24,7 @@ MAX_SEQUENCE_LENGTH = 512 -def train_tokenizer() -> Tuple[tokenizers.Tokenizer, Generator, int]: +def train_tokenizer() -> Tuple[tokenizers.Tokenizer, tf.data.Dataset]: tokenizer = tokenizers.Tokenizer(models.WordPiece(unk_token="")) tokenizer.decoder = decoders.WordPiece() @@ -39,14 +37,8 @@ def train_tokenizer() -> Tuple[tokenizers.Tokenizer, Generator, int]: pre_tokenizers.Whitespace(), pre_tokenizers.Digits(individual_digits=False) ]) - tokenizer.post_processor = processors.TemplateProcessing( - single="$A ", - pair="$A [SEP] $B:1", - special_tokens=[("[SEP]", 1), ("", 2), ("", 3)] - ) - # dataset = datasets.load_dataset("wikitext", "wikitext-103-raw-v1", split="test") - dataset = datasets.load_dataset("multi_news", split="test") + dataset = datasets.load_dataset("multi_news", split="train") def batch_iterator(batch_size=1000): for i in range(0, len(dataset), batch_size): @@ -95,8 +87,6 @@ def main(save_path: str): os.mkdir(save_path) tokenizer, data = train_tokenizer() - tokenizer.enable_truncation(2 * 512 + 1) # encoding + decoding + [SEP] token - with open(f"{save_path}/metadata.tsv", "w") as tsv: for token, _ in sorted(tokenizer.get_vocab().items(), key=lambda s: s[-1]): tsv.write(f"{token}\n") @@ -105,8 +95,8 @@ def main(save_path: str): max_sequence_len=MAX_SEQUENCE_LENGTH, embedding_size=EMBEDDING_SIZE, token_hash=tokenizer.get_vocab(), - input_drop_rate=0.3, - hidden_drop_rate=0.5 + input_drop_rate=0.2, + hidden_drop_rate=0.3 ) scheduler = RampUpDecaySchedule(EMBEDDING_SIZE, warmup_steps=4000) diff --git a/text2vec/autoencoders.py b/text2vec/autoencoders.py index 81c9295..0a5d0d7 100644 --- a/text2vec/autoencoders.py +++ b/text2vec/autoencoders.py @@ -89,43 +89,42 @@ def __init__(self, max_sequence_len: int, embedding_size: int, self.encode_layer = TransformerEncoder(n_stacks=1, num_layers=8, **params) self.decode_layer = TransformerDecoder(n_stacks=1, num_layers=8, **params) - def call(self, tokens, training: bool = False): # pylint: disable=missing-function-docstring - tokens = self.tokenizer(tokens) - x_enc, enc_mask, _ = self.embed_layer(tokens, training=training) - x_enc, context = self.encode_layer(x_enc, mask=enc_mask, training=training) - return x_enc, context, enc_mask + def call(self, inputs, training: bool = False): # pylint: disable=missing-function-docstring + encoding_text = inputs[0] + decoding_text = inputs[1] if len(inputs) > 1 else encoding_text + + encode_tokens = self.tokenizer(encoding_text) + x_embed, mask_encode, _ = self.embed_layer(encode_tokens, training=training) + x_encode, context = self.encode_layer(x_embed, mask=mask_encode, training=training) + + decode_tokens = self.tokenizer(decoding_text) + x_decode, mask_decode, _ = self.embed_layer(decode_tokens[:, :-1]) # skip + x_decode = self.decode_layer( + x_enc=x_encode, + x_dec=x_decode, + dec_mask=mask_decode, + context=context, + attention=self.encode_layer.attention, + training=training + ) - def train_step(self, data): # pylint: disable=missing-function-docstring - encoding_tok, decoding_tok = data - decoding_tok = self.tokenizer(decoding_tok) + return x_embed, x_decode, mask_decode, decode_tokens + def train_step(self, data): # pylint: disable=missing-function-docstring with tf.GradientTape() as tape: - with tf.name_scope('Encoding'): - x_enc, context, _ = self(encoding_tok, training=True) - - with tf.name_scope('Decoding'): - targets = decoding_tok[:, 1:] # skip the token with the slice on axis=1 - if isinstance(self.embed_layer, TokenEmbed): - targets = tf.ragged.map_flat_values(self.embed_layer.table.lookup, targets) - targets = self.embed_layer.slicer(targets) + _, x_decode, mask_decode, decode_tokens = self(data, training=True) - decoding_tok, dec_mask, _ = self.embed_layer(decoding_tok[:, :-1]) # skip - decoding_tok = self.decode_layer( - x_enc=x_enc, - x_dec=decoding_tok, - dec_mask=dec_mask, - context=context, - attention=self.encode_layer.attention, - training=True - ) - decoding_tok = tf.tensordot(decoding_tok, self.embed_layer.embeddings, axes=[2, 1]) + targets = decode_tokens[:, 1:] # skip the token with the slice on axis=1 + if isinstance(self.embed_layer, TokenEmbed): + targets = tf.ragged.map_flat_values(self.embed_layer.table.lookup, targets) + targets = self.embed_layer.slicer(targets) - loss = loss = tf.nn.sparse_softmax_cross_entropy_with_logits( - logits=decoding_tok, + loss = tf.nn.sparse_softmax_cross_entropy_with_logits( + logits=tf.tensordot(x_decode, self.embed_layer.embeddings, axes=[2, 1]), labels=targets.to_tensor(default_value=0) ) - loss = loss * dec_mask - loss = tf.math.reduce_sum(loss, axis=1) + loss = loss * mask_decode + # loss = tf.math.reduce_sum(loss, axis=1) loss = tf.reduce_mean(loss) gradients = tape.gradient(loss, self.trainable_variables) @@ -136,35 +135,21 @@ def train_step(self, data): # pylint: disable=missing-function-docstring return {"loss": loss, 'learning_rate': self.optimizer.learning_rate} def test_step(self, data): # pylint: disable=missing-function-docstring - encoding_tok, decoding_tok = data - decoding_tok = self.tokenizer(decoding_tok) - - with tf.name_scope('Encoding'): - x_enc, context, enc_mask = self(encoding_tok, training=False) + _, x_decode, mask_decode, decode_tokens = self(data, training=False) - with tf.name_scope('Decoding'): - targets = tf.ragged.map_flat_values(self.embed_layer.table.lookup, decoding_tok[:, 1:]) # skip - targets = self.embed_layer.slicer(targets) - - decoding_tok, dec_mask, _ = self.embed_layer(decoding_tok[:, :-1]) # skip - decoding_tok = self.decode_layer( - x_enc=x_enc, - enc_mask=enc_mask, - x_dec=decoding_tok, - dec_mask=dec_mask, - context=context, - attention=self.encode_layer.attention, - training=False - ) - decoding_tok = tf.tensordot(decoding_tok, self.embed_layer.embeddings, axes=[2, 1]) + targets = decode_tokens[:, 1:] # skip the token with the slice on axis=1 + if isinstance(self.embed_layer, TokenEmbed): + targets = tf.ragged.map_flat_values(self.embed_layer.table.lookup, targets) + targets = self.embed_layer.slicer(targets) loss = tf.nn.sparse_softmax_cross_entropy_with_logits( - logits=decoding_tok, + logits=tf.tensordot(x_decode, self.embed_layer.embeddings, axes=[2, 1]), labels=targets.to_tensor(default_value=0) ) - loss = loss * dec_mask - loss = tf.math.reduce_sum(loss, axis=1) + loss = loss * mask_decode + # loss = tf.math.reduce_sum(loss, axis=1) loss = tf.reduce_mean(loss) + return {"loss": loss, **{m.name: m.result() for m in self.metrics}} @tf.function(input_signature=[tf.TensorSpec(shape=[None], dtype=tf.string)]) @@ -183,11 +168,13 @@ def embed(self, sentences) -> Dict[str, tf.Tensor]: and (batch_size, max_sequence_len, embedding_size) respectively. """ - sequences, attention, _ = self(sentences, training=False) - return {"sequences": sequences, "attention": attention} + tokens = self.tokenizer(sentences) + x, mask, _ = self.embed_layer(tokens, training=False) + x, context = self.encode_layer(x, mask=mask, training=False) + return {"sequences": x, "attention": context} @tf.function(input_signature=[tf.TensorSpec(shape=[None], dtype=tf.string)]) - def token_embed(self, sentences) -> Dict[str, tf.Tensor]: + def token_embed(self, sentences) -> Dict[str, tf.RaggedTensor]: """Takes batches of free text and returns word embeddings along with the associate token. Parameters @@ -197,16 +184,13 @@ def token_embed(self, sentences) -> Dict[str, tf.Tensor]: Returns ------- - Dict[str, tf.Tensor] - Padded tokens and embedding vectors with shapes (batch_size, max_sequence_len) - and (batch_size, max_sequence_len, embedding_size) respectively. + Dict[str, tf.RaggedTensor] + Ragged tokens and embedding tensors with shapes (batch_size, None) + and (batch_size, None, embedding_size) respectively. """ tokens = self.tokenizer(sentences) - return { - "tokens": tokens.to_tensor(''), - "embeddings": self.embed_layer.get_embedding(tokens).to_tensor(0) - } + return {"tokens": tokens, "embeddings": self.embed_layer.get_embedding(tokens)} class LstmAutoEncoder(Model): diff --git a/text2vec/models/components/attention.py b/text2vec/models/components/attention.py index 922821b..50400ec 100644 --- a/text2vec/models/components/attention.py +++ b/text2vec/models/components/attention.py @@ -38,20 +38,20 @@ class ScaledDotAttention(layers.Layer): def __init__(self): super().__init__(name="ScaledDotAttention") - self.scale = tf.constant(1e9, dtype=tf.float32) + self.neg_inf = tf.constant(-1e9, dtype=tf.float32) # pylint: disable=missing-function-docstring def call(self, query, key, value, mask_future: bool = False): - numerator = tf.einsum('ijk,ilk->ijl', query, key) + numerator = tf.matmul(query, key, transpose_b=True) denominator = tf.sqrt(tf.cast(tf.shape(key)[-1], tf.float32)) if mask_future: - upper = (1 + self.scale) * tf.linalg.band_part(tf.ones_like(numerator), num_lower=0, num_upper=-1) - mask = 1 - upper - numerator *= mask + upper = tf.linalg.band_part(tf.ones(tf.shape(numerator)[1:], dtype=tf.float32), num_lower=0, num_upper=-1) + diag = tf.linalg.band_part(upper, num_lower=0, num_upper=0) + numerator += (self.neg_inf * (upper - diag)) x = tf.nn.softmax(numerator / denominator) - return tf.einsum('ijk,ikl->ijl', x, value) + return tf.matmul(x, value) class BahdanauAttention(layers.Layer): diff --git a/text2vec/models/components/text_inputs.py b/text2vec/models/components/text_inputs.py index 888ef4b..22c189c 100644 --- a/text2vec/models/components/text_inputs.py +++ b/text2vec/models/components/text_inputs.py @@ -48,19 +48,22 @@ def __init__(self, vocab_size: int, embedding_size: int, max_sequence_len: int): dtype=tf.float32, trainable=True ) + self.sqrt_d = tf.math.sqrt(tf.cast(embedding_size, tf.float32)) self.max_len = tf.constant(max_sequence_len) self.slicer = layers.Lambda(lambda x: x[:, :max_sequence_len], name="sequence-slice") def call(self, token_ids, **kwargs): token_ids = self.slicer(token_ids) x = tf.ragged.map_flat_values(tf.nn.embedding_lookup, self.embeddings, token_ids) + x * self.sqrt_d x = x.to_tensor(0) - x = x * tf.math.sqrt(tf.cast(tf.shape(self.embeddings)[-1], tf.float32)) # sqrt(embedding_size) - seq_lengths = token_ids.row_lengths() - time_steps = tf.cast(tf.reduce_max(seq_lengths), tf.int32) - mask = tf.sequence_mask(lengths=seq_lengths, maxlen=time_steps, dtype=tf.float32) - return x, mask, time_steps + mask = tf.sequence_mask( + lengths=token_ids.row_lengths(), + maxlen=token_ids.bounding_shape()[-1], + dtype=tf.float32 + ) + return x, mask, token_ids.row_lengths() def get_embedding(self, token_ids: tf.RaggedTensor) -> tf.RaggedTensor: """Get the token embeddings for the input IDs. diff --git a/text2vec/models/components/utils.py b/text2vec/models/components/utils.py index 9da7012..c9a6424 100644 --- a/text2vec/models/components/utils.py +++ b/text2vec/models/components/utils.py @@ -1,6 +1,6 @@ import numpy as np import tensorflow as tf -from tensorflow.keras import layers +from tensorflow.keras import layers, initializers class LayerNorm(layers.Layer): @@ -63,13 +63,9 @@ def __init__(self): super().__init__(name="TensorProjection") def call(self, x, projection_vector): + projection_vector = tf.math.l2_normalize(projection_vector, axis=-1) inner_product = tf.einsum("ijk,ik->ij", x, projection_vector) - time_steps = tf.shape(x)[1] - p_vector_norm_squared = tf.norm(projection_vector, axis=1) ** 2 - p_vector_norm_squared = tf.tile(tf.expand_dims(p_vector_norm_squared, -1), [1, time_steps]) - - alpha = tf.divide(inner_product, p_vector_norm_squared) - return tf.einsum("ij,ik->ijk", alpha, projection_vector) + return tf.einsum("ij,ik->ijk", inner_product, projection_vector) class PositionalEncoder(layers.Layer): @@ -120,4 +116,50 @@ def __init__(self, emb_dims: int, max_sequence_len: int): def call(self, x: tf.Tensor, mask: tf.Tensor): time_steps = tf.shape(x)[1] - return tf.einsum('ijk,ij->ijk', x + self.encoder[:time_steps, :], mask) + return tf.expand_dims(mask, axis=-1) * (x + self.encoder[:time_steps, ...]) + + +class VariationPositionalEncoder(layers.Layer): + """Learns the relative phases between sequence steps in an attention-based transformer, where there is no + inherent sequential ordering. + + Parameters + ---------- + emb_dims : int + The word-embedding dimensionality. This value determines the dimensionalities of the hidden weights. + max_sequence_len : int + Longest sequence seen at training time. + + Examples + -------- + ```python + import tensorflow as tf + from text2vec.models import TextInput + from text2vec.models import utils + + lookup = {'string': 0, 'is': 1, 'example': 2} + inputer = TextInput(token_hash=lookup, embedding_size=16, max_sequence_len=10) + encoder = utils.VariationPositionalEncoder(emb_dims=16, max_sequence_len=10) + + text = tf.ragged.constant([ + ["Sample", "string", "."], + ["This", "is", "a", "second", "example", "."] + ]) + x, mask, _ = inputer(text) + encoder(x, mask) + ``` + """ + + def __init__(self, emb_dims: int, max_sequence_len: int): + super().__init__() + + self.encoder = tf.Variable( + initializers.GlorotUniform()(shape=(max_sequence_len, emb_dims)), + dtype=tf.float32, + trainable=True, + name="positional-encoder" + ) + + def call(self, x: tf.Tensor, mask: tf.Tensor): + time_steps = tf.shape(x)[1] + return tf.expand_dims(mask, axis=-1) * (x + self.encoder[:time_steps, ...]) diff --git a/text2vec/models/transformer.py b/text2vec/models/transformer.py index 5546c66..6d060f6 100644 --- a/text2vec/models/transformer.py +++ b/text2vec/models/transformer.py @@ -3,7 +3,7 @@ from .components.attention import BahdanauAttention from .components.attention import MultiHeadAttention from .components.feed_forward import PositionWiseFFN -from .components.utils import PositionalEncoder +from .components.utils import VariationPositionalEncoder from .components.utils import LayerNorm from .components.utils import TensorProjection @@ -51,17 +51,17 @@ def __init__(self, max_sequence_len, num_layers=8, n_stacks=1, embedding_size=50 super().__init__() dims = embedding_size - self.drop = layers.Dropout(input_drop_rate) - self.h_drop = layers.Dropout(hidden_drop_rate) + self.positional_encode = VariationPositionalEncoder(emb_dims=dims, max_sequence_len=max_sequence_len) self.layer_norm = LayerNorm() - - self.positional_encode = PositionalEncoder(emb_dims=dims, max_sequence_len=max_sequence_len) self.MHA = [ MultiHeadAttention(emb_dims=dims, num_layers=num_layers, drop_rate=input_drop_rate) for _ in range(n_stacks) ] self.FFN = [PositionWiseFFN(emb_dims=dims) for _ in range(n_stacks)] - self.attention = BahdanauAttention(size=dims) + self.attention = BahdanauAttention(size=dims, drop_rate=hidden_drop_rate) + + self.drop = layers.Dropout(input_drop_rate) + self.h_drop = layers.Dropout(hidden_drop_rate) # pylint: disable=missing-function-docstring def call(self, x, mask, training: bool = False): @@ -102,18 +102,18 @@ def __init__(self, max_sequence_len, num_layers=8, n_stacks=1, embedding_size=50 super().__init__() dims = embedding_size - self.drop = layers.Dropout(input_drop_rate) - self.h_drop = layers.Dropout(hidden_drop_rate) self.layer_norm = LayerNorm() self.projection = TensorProjection() - - self.positional_encode = PositionalEncoder(emb_dims=dims, max_sequence_len=max_sequence_len) + self.positional_encode = VariationPositionalEncoder(emb_dims=dims, max_sequence_len=max_sequence_len) self.MHA = [ MultiHeadAttention(emb_dims=dims, num_layers=num_layers, drop_rate=input_drop_rate) for _ in range(n_stacks) ] self.FFN = [PositionWiseFFN(emb_dims=dims) for _ in range(n_stacks)] + self.drop = layers.Dropout(input_drop_rate) + self.h_drop = layers.Dropout(hidden_drop_rate) + # pylint: disable=missing-function-docstring def call(self, x_enc, x_dec, dec_mask, context, attention: BahdanauAttention, training: bool = False): x_dec = self.positional_encode(x_dec, dec_mask) From d7aba55afba91b3ac4f9f7616ccca81378da76d1 Mon Sep 17 00:00:00 2001 From: Dave Date: Thu, 7 Jul 2022 12:09:39 -0400 Subject: [PATCH 06/10] sync up lstm autoencoders with transformers --- .gitignore | 2 +- .../trainers/{wiki_lstm.py => news_lstm.py} | 98 +++++++------- examples/trainers/news_transformer.py | 3 +- setup.py | 2 +- text2vec/autoencoders.py | 123 ++++++++---------- text2vec/models/sequential.py | 11 +- 6 files changed, 112 insertions(+), 127 deletions(-) rename examples/trainers/{wiki_lstm.py => news_lstm.py} (55%) diff --git a/.gitignore b/.gitignore index 55960d9..8456922 100644 --- a/.gitignore +++ b/.gitignore @@ -7,7 +7,7 @@ **/*.dev.yml wiki_t2v/ -multi_news_t2v/ +multi_news_t2v*/ # JavaScript configs and dependencies **/.eslintrc.json diff --git a/examples/trainers/wiki_lstm.py b/examples/trainers/news_lstm.py similarity index 55% rename from examples/trainers/wiki_lstm.py rename to examples/trainers/news_lstm.py index faffaf2..5b136ba 100644 --- a/examples/trainers/wiki_lstm.py +++ b/examples/trainers/news_lstm.py @@ -1,4 +1,4 @@ -from typing import Generator, List, Tuple, Union +from typing import Tuple import os import datasets @@ -7,22 +7,23 @@ from tokenizers import decoders from tokenizers import normalizers from tokenizers import pre_tokenizers -from tokenizers import processors from tokenizers import trainers -from nltk.tokenize import PunktSentenceTokenizer -import numpy as np import tensorflow as tf +from tensorflow.keras import optimizers, callbacks +from tensorflow.keras import backend as K from tensorboard.plugins import projector from text2vec.autoencoders import LstmAutoEncoder from text2vec.optimizer_tools import RampUpDecaySchedule os.environ["TOKENIZERS_PARALLELISM"] = "true" -sent_tokenizer = PunktSentenceTokenizer().tokenize +root = os.path.dirname(os.path.abspath(__file__)) +EMBEDDING_SIZE = 128 +MAX_SEQUENCE_LENGTH = 512 -def train_tokenizer() -> Tuple[tokenizers.Tokenizer, Generator, int]: +def train_tokenizer() -> Tuple[tokenizers.Tokenizer, tf.data.Dataset]: tokenizer = tokenizers.Tokenizer(models.WordPiece(unk_token="")) tokenizer.decoder = decoders.WordPiece() @@ -35,17 +36,13 @@ def train_tokenizer() -> Tuple[tokenizers.Tokenizer, Generator, int]: pre_tokenizers.Whitespace(), pre_tokenizers.Digits(individual_digits=False) ]) - tokenizer.post_processor = processors.TemplateProcessing( - single="$A ", - pair="$A [SEP] $B:1", - special_tokens=[("[SEP]", 1), ("", 2), ("", 3)] - ) - dataset = datasets.load_dataset("wikitext", "wikitext-103-raw-v1", split="test") + dataset = datasets.load_dataset("multi_news", split="train") def batch_iterator(batch_size=1000): for i in range(0, len(dataset), batch_size): - yield dataset[i: i + batch_size]["text"] + for key in dataset.features: + yield dataset[i: i + batch_size][key] tokenizer.train_from_iterator( batch_iterator(), @@ -55,14 +52,32 @@ def batch_iterator(batch_size=1000): ) ) + tokenizer.enable_truncation(2 * MAX_SEQUENCE_LENGTH + 3) # 2 for the [SEP], , tokens + tokenizer.post_processor = tokenizers.processors.TemplateProcessing( + single="$A", + pair="$A:0 [SEP] $B:1 ", + special_tokens=[ + ("[SEP]", 1), + ("", 2), + ("", 3) + ] + ) + def generator(): for record in dataset: - if record['text'].strip() != '': - for sentence in sent_tokenizer(record['text']): - yield sentence - - data = tf.data.Dataset.from_generator(generator, output_signature=(tf.TensorSpec(shape=(None), dtype=tf.string))) - data = data.map(tf.strings.strip, num_parallel_calls=tf.data.experimental.AUTOTUNE) + if record["document"] and record["summary"]: + enc, dec = ' '.join(tokenizer.encode( + record["document"], + pair=record["summary"] + ).tokens).split(' [SEP] ', maxsplit=2) + + if enc.strip() != "" and dec != "": + yield enc, dec + + data = tf.data.Dataset.from_generator( + generator, + output_signature=(tf.TensorSpec(shape=(None), dtype=tf.string), tf.TensorSpec(shape=(None), dtype=tf.string)) + ) return tokenizer, data @@ -71,38 +86,20 @@ def main(save_path: str): os.mkdir(save_path) tokenizer, data = train_tokenizer() - tokenizer.enable_truncation(2 * 512 + 1) # encoding + decoding + [SEP] token - with open(f"{save_path}/metadata.tsv", "w") as tsv: for token, _ in sorted(tokenizer.get_vocab().items(), key=lambda s: s[-1]): tsv.write(f"{token}\n") - def encode(x): - def token_mapper(text: Union[str, List[str]]): - text = text.numpy() - - if isinstance(text, np.ndarray): - enc, dec = [], [] - for batch in tokenizer.encode_batch([(t.decode('utf8'), t.decode('utf8')) for t in text]): - enc_, dec_ = ' '.join(batch.tokens).split(' [SEP] ') - enc.append(enc_) - dec.append(dec_) - return (enc, dec) - - text = text.decode('utf8') - enc, dec = ' '.join(tokenizer.encode(text, pair=text).tokens).split(' [SEP] ') - return (enc, dec) - - return tf.py_function(token_mapper, inp=[x], Tout=[tf.string, tf.string]) - model = LstmAutoEncoder( - max_sequence_len=512, - embedding_size=128, + max_sequence_len=MAX_SEQUENCE_LENGTH, + embedding_size=EMBEDDING_SIZE, token_hash=tokenizer.get_vocab(), input_drop_rate=0.3, hidden_drop_rate=0.5 ) - model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=RampUpDecaySchedule(embedding_size=128))) + + scheduler = RampUpDecaySchedule(EMBEDDING_SIZE, warmup_steps=4000) + model.compile(optimizer=optimizers.Adam(scheduler(0).numpy())) checkpoint = tf.train.Checkpoint(Classifier=model, optimizer=model.optimizer) checkpoint_manager = tf.train.CheckpointManager(checkpoint, save_path, max_to_keep=3) @@ -117,16 +114,17 @@ def token_mapper(text: Union[str, List[str]]): embeddings_config.metadata_path = f"{save_path}/metadata.tsv" projector.visualize_embeddings(logdir=save_path, config=config) - data = data.map(encode, num_parallel_calls=tf.data.experimental.AUTOTUNE) model.fit( x=data.prefetch(8).batch(64), callbacks=[ - tf.keras.callbacks.TensorBoard( - log_dir=save_path, - write_graph=True, - update_freq=100 - ), - tf.keras.callbacks.LambdaCallback(on_epoch_end=lambda epoch, logs: checkpoint_manager.save()) + callbacks.TensorBoard(log_dir=save_path, write_graph=True, update_freq=100), + callbacks.LambdaCallback( + on_epoch_end=lambda epoch, logs: checkpoint_manager.save(), + on_batch_end=lambda batch, logs: K.set_value( + model.optimizer.lr, + K.get_value(scheduler(model.optimizer.iterations)) + ) + ) ], epochs=1 ) @@ -142,4 +140,4 @@ def token_mapper(text: Union[str, List[str]]): if __name__ == '__main__': - main(save_path='./wiki_t2v') + main(save_path=f'{root}/../../multi_news_t2v_sequential') diff --git a/examples/trainers/news_transformer.py b/examples/trainers/news_transformer.py index 41408d6..5d8777b 100644 --- a/examples/trainers/news_transformer.py +++ b/examples/trainers/news_transformer.py @@ -1,4 +1,4 @@ -from typing import Generator, Tuple +from typing import Tuple import os import datasets @@ -7,7 +7,6 @@ from tokenizers import decoders from tokenizers import normalizers from tokenizers import pre_tokenizers -from tokenizers import processors from tokenizers import trainers import tensorflow as tf diff --git a/setup.py b/setup.py index 37d2ef0..b8ad7ee 100644 --- a/setup.py +++ b/setup.py @@ -3,7 +3,7 @@ setup( name="text2vec", - version="1.4.0", + version="2.0.0", description="Building blocks for text vectorization and embedding", author="Dave Hollander", author_url="https://github.com/brainsqueeze", diff --git a/text2vec/autoencoders.py b/text2vec/autoencoders.py index 0a5d0d7..f520bdd 100644 --- a/text2vec/autoencoders.py +++ b/text2vec/autoencoders.py @@ -271,44 +271,48 @@ def __init__(self, max_sequence_len: int, embedding_size: int, num_hidden: int = self.encode_layer = RecurrentEncoder(num_hidden=num_hidden, **params) self.decode_layer = RecurrentDecoder(num_hidden=num_hidden, **params) - def call(self, tokens, training: bool = False): # pylint: disable=missing-function-docstring - tokens = self.tokenizer(tokens) - x_enc, enc_mask, _ = self.embed_layer(tokens, training=training) - x_enc, context, *states = self.encode_layer(x_enc, mask=enc_mask, training=training) - return x_enc, context, enc_mask, states + def call(self, inputs, training: bool = False): # pylint: disable=missing-function-docstring + # tokens = self.tokenizer(tokens) + # x_enc, enc_mask, _ = self.embed_layer(tokens, training=training) + # x_enc, context, *states = self.encode_layer(x_enc, mask=enc_mask, training=training) + # return x_enc, context, enc_mask, states - def train_step(self, data): # pylint: disable=missing-function-docstring - encoding_tok, decoding_tok = data - decoding_tok = self.tokenizer(decoding_tok) + encoding_text = inputs[0] + decoding_text = inputs[1] if len(inputs) > 1 else encoding_text + + encode_tokens = self.tokenizer(encoding_text) + x_embed, mask_encode, _ = self.embed_layer(encode_tokens, training=training) + x_encode, context, *states = self.encode_layer(x_embed, mask=mask_encode, training=training) + + decode_tokens = self.tokenizer(decoding_text) + x_decode, mask_decode, _ = self.embed_layer(decode_tokens[:, :-1]) # skip + x_decode = self.decode_layer( + x_enc=x_encode, + x_dec=x_decode, + dec_mask=mask_decode, + context=context, + # attention=self.encode_layer.attention, + initial_state=states, + training=training + ) + return x_embed, x_decode, mask_decode, decode_tokens + + def train_step(self, data): # pylint: disable=missing-function-docstring with tf.GradientTape() as tape: - with tf.name_scope('Encoding'): - x_enc, context, enc_mask, states = self(encoding_tok, training=True) - - with tf.name_scope('Decoding'): - targets = decoding_tok[:, 1:] # skip the token with the slice on axis=1 - if isinstance(self.embed_layer, TokenEmbed): - targets = tf.ragged.map_flat_values(self.embed_layer.table.lookup, targets) - targets = self.embed_layer.slicer(targets) - - decoding_tok, dec_mask, _ = self.embed_layer(decoding_tok[:, :-1]) - decoding_tok = self.decode_layer( - x_enc=x_enc, - enc_mask=enc_mask, - x_dec=decoding_tok, - dec_mask=dec_mask, - context=context, - initial_state=states, - training=True - ) - decoding_tok = tf.tensordot(decoding_tok, self.embed_layer.embeddings, axes=[2, 1]) - - loss = loss = tf.nn.sparse_softmax_cross_entropy_with_logits( - logits=decoding_tok, + _, x_decode, mask_decode, decode_tokens = self(data, training=True) + + targets = decode_tokens[:, 1:] # skip the token with the slice on axis=1 + if isinstance(self.embed_layer, TokenEmbed): + targets = tf.ragged.map_flat_values(self.embed_layer.table.lookup, targets) + targets = self.embed_layer.slicer(targets) + + loss = tf.nn.sparse_softmax_cross_entropy_with_logits( + logits=tf.tensordot(x_decode, self.embed_layer.embeddings, axes=[2, 1]), labels=targets.to_tensor(default_value=0) ) - loss = loss * dec_mask - loss = tf.math.reduce_sum(loss, axis=1) + loss = loss * mask_decode + # loss = tf.math.reduce_sum(loss, axis=1) loss = tf.reduce_mean(loss) gradients = tape.gradient(loss, self.trainable_variables) @@ -319,36 +323,19 @@ def train_step(self, data): # pylint: disable=missing-function-docstring return {"loss": loss, 'learning_rate': self.optimizer.learning_rate} def test_step(self, data): # pylint: disable=missing-function-docstring - encoding_tok, decoding_tok = data - decoding_tok = self.tokenizer(decoding_tok) - - with tf.name_scope('Encoding'): - x_enc, context, enc_mask, states = self(encoding_tok, training=False) - - with tf.name_scope('Decoding'): - targets = decoding_tok[:, 1:] # skip the token with the slice on axis=1 - if isinstance(self.embed_layer, TokenEmbed): - targets = tf.ragged.map_flat_values(self.embed_layer.table.lookup, targets) - targets = self.embed_layer.slicer(targets) + _, x_decode, mask_decode, decode_tokens = self(data, training=False) - decoding_tok, dec_mask, _ = self.embed_layer(decoding_tok[:, :-1]) - decoding_tok = self.decode_layer( - x_enc=x_enc, - enc_mask=enc_mask, - x_dec=decoding_tok, - dec_mask=dec_mask, - context=context, - initial_state=states, - training=False - ) - decoding_tok = tf.tensordot(decoding_tok, self.embed_layer.embeddings, axes=[2, 1]) + targets = decode_tokens[:, 1:] # skip the token with the slice on axis=1 + if isinstance(self.embed_layer, TokenEmbed): + targets = tf.ragged.map_flat_values(self.embed_layer.table.lookup, targets) + targets = self.embed_layer.slicer(targets) - loss = loss = tf.nn.sparse_softmax_cross_entropy_with_logits( - logits=decoding_tok, + loss = tf.nn.sparse_softmax_cross_entropy_with_logits( + logits=tf.tensordot(x_decode, self.embed_layer.embeddings, axes=[2, 1]), labels=targets.to_tensor(default_value=0) ) - loss = loss * dec_mask - loss = tf.math.reduce_sum(loss, axis=1) + loss = loss * mask_decode + # loss = tf.math.reduce_sum(loss, axis=1) loss = tf.reduce_mean(loss) return {"loss": loss, **{m.name: m.result() for m in self.metrics}} @@ -369,11 +356,13 @@ def embed(self, sentences) -> Dict[str, tf.Tensor]: and (batch_size, max_sequence_len, embedding_size) respectively. """ - sequences, attention, *_ = self(sentences, training=False) - return {"sequences": sequences, "attention": attention} + tokens = self.tokenizer(sentences) + x, mask, _ = self.embed_layer(tokens, training=False) + x, context, *_ = self.encode_layer(x, mask=mask, training=False) + return {"sequences": x, "attention": context} @tf.function(input_signature=[tf.TensorSpec(shape=[None], dtype=tf.string)]) - def token_embed(self, sentences) -> Dict[str, tf.Tensor]: + def token_embed(self, sentences) -> Dict[str, tf.RaggedTensor]: """Takes batches of free text and returns word embeddings along with the associate token. Parameters @@ -383,13 +372,13 @@ def token_embed(self, sentences) -> Dict[str, tf.Tensor]: Returns ------- - Dict[str, tf.Tensor] - Padded tokens and embedding vectors with shapes (batch_size, max_sequence_len) - and (batch_size, max_sequence_len, embedding_size) respectively. + Dict[str, tf.RaggedTensor] + Ragged tokens and embedding tensors with shapes (batch_size, None) + and (batch_size, None, embedding_size) respectively. """ tokens = self.tokenizer(sentences) return { - "tokens": tokens.to_tensor(''), - "embeddings": self.embed_layer.get_embedding(tokens).to_tensor(0) + "tokens": tokens, + "embeddings": self.embed_layer.get_embedding(tokens) } diff --git a/text2vec/models/sequential.py b/text2vec/models/sequential.py index 1f533ea..0e1a964 100644 --- a/text2vec/models/sequential.py +++ b/text2vec/models/sequential.py @@ -40,13 +40,14 @@ class RecurrentEncoder(layers.Layer): ``` """ - def __init__(self, max_sequence_len, num_hidden, num_layers=2, input_drop_rate: float = 0., **kwargs): + def __init__(self, max_sequence_len, num_hidden, num_layers=2, + input_drop_rate: float = 0., hidden_drop_rate: float = 0., **kwargs): super().__init__() self.max_sequence_length = max_sequence_len self.drop = layers.Dropout(input_drop_rate) self.bi_lstm = BidirectionalLSTM(num_layers=num_layers, num_hidden=num_hidden, return_states=True) - self.attention = BahdanauAttention(size=2 * num_hidden) + self.attention = BahdanauAttention(size=2 * num_hidden, drop_rate=hidden_drop_rate) # pylint: disable=missing-function-docstring def call(self, x, mask, training: bool = False): @@ -54,7 +55,7 @@ def call(self, x, mask, training: bool = False): mask = tf.expand_dims(mask, axis=-1) x = self.drop(x, training=training) x, states = self.bi_lstm(x * mask, training=training) - context = self.attention(x * mask) + x, context = self.attention(x * mask) if training: return x, context, states @@ -94,11 +95,9 @@ def __init__(self, max_sequence_len, num_hidden, embedding_size=50, num_layers=2 self.dense = layers.Dense(units=dims, activation=tf.nn.relu) # pylint: disable=missing-function-docstring - def call(self, x_enc, enc_mask, x_dec, dec_mask, context, training: bool = False): - enc_mask = tf.expand_dims(enc_mask, axis=-1) + def call(self, x_enc, x_dec, dec_mask, context, initial_state=None, training: bool = False): dec_mask = tf.expand_dims(dec_mask, axis=-1) - initial_state = kwargs.get("initial_state") x = self.drop(x_dec * dec_mask, training=training) if initial_state is not None: x = self.bi_lstm(x * dec_mask, initial_states=initial_state[0], training=training) From 0770fda1121b23070ff672158fa0166c5fb2a28d Mon Sep 17 00:00:00 2001 From: Dave Date: Thu, 7 Jul 2022 13:32:21 -0400 Subject: [PATCH 07/10] Remove deprecated references --- README.md | 145 +++++---------------- setup.py | 5 - text2vec/autoencoders.py | 19 +-- text2vec/models/__init__.py | 3 +- text2vec/models/components/feed_forward.py | 3 +- text2vec/models/components/feeder.py | 115 ---------------- text2vec/models/components/text_inputs.py | 31 +++++ text2vec/models/components/utils.py | 8 +- text2vec/models/sequential.py | 4 +- text2vec/models/transformer.py | 11 +- 10 files changed, 75 insertions(+), 269 deletions(-) delete mode 100644 text2vec/models/components/feeder.py diff --git a/README.md b/README.md index 0fed1ea..e60666f 100644 --- a/README.md +++ b/README.md @@ -27,119 +27,32 @@ pip install git+https://github.com/brainsqueeze/text2vec.git ## Motivation --- -Word embedding models have been very beneficial to natural -language processing. The technique is able to distill semantic -meaning from words by treating them as vectors in a -high-dimensional vector space. - -This package attempts to accomplish the same semantic embedding, -but do this at the sentence and paragraph level. Within a -sentence the order of words and the use of punctuation and -conjugations are very important for extracting the meaning -of blocks of text. - -Inspiration is taken from recent advances in text summary -models (pointer-generator), where an attention mechanism -[[1](https://arxiv.org/abs/1409.0473)] is -used to extrude the overall meaning of the input text. In the -case of text2vec, we use the attention vectors found from the -input text as the embedded vector representing the input. -Furthermore, recent attention-only approaches to sequence-to-sequence -modeling are adapted. - -**note**: this is not a canonical implementation of the attention -mechanism, but this method was chosen intentionally to be able to -leverage the attention vector as the embedding output. +Word embedding models have been very beneficial to natural language processing. The technique is able to distill semantic meaning from words by treating them as vectors in a high-dimensional vector space. + +This package attempts to accomplish the same semantic embedding, but do this at the sentence and paragraph level. Within a sentence the order of words and the use of punctuation and conjugations are very important for extracting the meaning of blocks of text. + +Inspiration is taken from recent advances in text summary models (pointer-generator), where an attention mechanism [[1](https://arxiv.org/abs/1409.0473)] is used to extrude the overall meaning of the input text. In the case of text2vec, we use the attention vectors found from the input text as the embedded vector representing the input. Furthermore, recent attention-only approaches to sequence-to-sequence modeling are adapted. + ### Transformer model --- -This is a tensor-to-tensor model adapted from the work in -[Attention Is All You Need](https://arxiv.org/abs/1706.03762). -The embedding and encoding steps follow directly from -[[2](https://arxiv.org/abs/1706.03762)], however a self- -attention is applied at the end of the encoding steps and a -context-vector is learned, which in turn is used to project -the decoding tensors onto. - -The decoding steps begin as usual with the word-embedded input -sequences shifted right, then multi-head attention, skip connection -and layer-normalization is applied. Before continuing, we project -the resulting decoded sequences onto the context-vector from the -encoding steps. The projected tensors are then passed through -the position-wise feed-forward (conv1D) + skip connection and layer- -normalization again, before once more being projected onto the -context-vectors. +This is a tensor-to-tensor model adapted from the work in [Attention Is All You Need](https://arxiv.org/abs/1706.03762). The embedding and encoding steps follow directly from [[2](https://arxiv.org/abs/1706.03762)], however a self-attention is applied at the end of the encoding steps and a context-vector is learned, which in turn is used to project the decoding tensors onto. + +The decoding steps begin as usual with the word-embedded input sequences shifted right, then multi-head attention, skip connection and layer-normalization is applied. Before continuing, we project the resulting decoded sequences onto the context-vector from the encoding steps. The projected tensors are then passed through the position-wise feed-forward (conv1D) + skip connection and layer-normalization again, before once more being projected onto the context-vectors. ### LSTM seq2seq -This is an adapted bi-directional LSTM encoder-decoder model with -a self-attention mechanism learned from the encoding steps. The -context-vectors are used to project the resulting decoded sequences -before computing logits. +This is an adapted bi-directional LSTM encoder-decoder model with a self-attention mechanism learned from the encoding steps. The context-vectors are used to project the resulting decoded sequences before computing logits. ## Training --- -Both models are trained using Adam SGD with the learning-rate decay -program in [[2](https://arxiv.org/abs/1706.03762)]. - -The pre-built auto-encoder models inherit from [tf.keras.Model](https://www.tensorflow.org/api_docs/python/tf/keras/Model), and as such they can be trained using the `fit` method. -An example of training on Wikitext data is available in the [examples folder](./examples/trainers/wiki_transformer.py). This uses HuggingFace [tokenizers](https://huggingface.co/docs/tokenizers/python/latest/) and [datasets](https://huggingface.co/docs/datasets/master/). - -Training the LSTM model can be initiated with -```bash -text2vec_main --run=train --yaml_config=/path/to/config.yml -``` -The training configuration YAML for attention models must look like -```yaml -training: - tokens: 10000 - max_sequence_length: 512 - epochs: 100 - batch_size: 64 - -model: - name: transformer_test - parameters: - embedding: 128 - layers: 8 - storage_dir: /path/to/save/model -``` -The `parameters` for recurrent models must include at least -`embedding` and `hidden`, which referes to the dimensionality of the hidden LSTM layer. The `training` section of the YAML file can also include user-defined sentences to use as a context-angle evaluation set. This can look like -```yaml -eval_sentences: - - The movie was great! - - The movie was terrible. -``` -It can also include a `data` tag which is a list of absolute file paths for custom training data sets. This can look like -```yaml -data_files: - - ~/path/to/data/set1.txt - - ~/path/to/data/set2.txt - ... -``` - -Likewise, the transformer model can be trained with -```bash -text2vec_main --run=train --attention --yaml_config=/path/to/config.yml -``` - -To view the output of training you can then run -```bash -tensorboard --logdir text_embedding -``` - -If you have CUDA and cuDNN installed you can run -`pip install -r requirements-gpu.txt`. -The GPU will automatically be detected and used if present, otherwise -it will fall back to the CPU for training and inferencing. +Both models are trained using Adam SGD with the learning-rate decay program in [[2](https://arxiv.org/abs/1706.03762)]. -### Mutual contextual orthogonality +The pre-built auto-encoder models inherit from [tf.keras.Model](https://www.tensorflow.org/api_docs/python/tf/keras/Model), and as such they can be trained using the [fit method](https://www.tensorflow.org/api_docs/python/tf/keras/Model#fit). An example of training on Wikitext data is available in the [examples folder](./examples/trainers/wiki_transformer.py). This uses HuggingFace [tokenizers](https://huggingface.co/docs/tokenizers/python/latest/) and [datasets](https://huggingface.co/docs/datasets/master/). -To impose quasi-mutual orthogonality on the learned context vectors simply add the `--orthogonal` flag to the training command. This will add a loss term that can be thought of as a Lagrange multiplier where the constraint is self-alignment of the context vectors, and orthogonality between non-self vectors. The aim is not to impose orthogonality between all text inputs that are not the same, but rather to coerce the model to learn significantly different encodings for different contextual inputs. ## Python API @@ -149,33 +62,32 @@ Text2vec includes a Python API with convenient classes for handling attention an #### Auto-encoders - - [text2vec.autoencoders.TransformerAutoEncoder](/text2vec/autoencoders.py#L13) - - [text2vec.autoencoders.LstmAutoEncoder](/text2vec/models/transformer.py#L134) + - [text2vec.autoencoders.TransformerAutoEncoder](/text2vec/autoencoders.py#L12) + - [text2vec.autoencoders.LstmAutoEncoder](/text2vec/models/transformer.py#L190) #### Layers - - [text2vec.models.TransformerEncoder](/text2vec/models/transformer.py#L11) - - [text2vec.models.TransformerDecoder](/text2vec/models/transformer.py#L81) - - [text2vec.models.RecurrentEncoder](/text2vec/models/sequential.py#L8) - - [text2vec.models.RecurrentDecoder](/text2vec/models/sequential.py#L61) + - [text2vec.models.TransformerEncoder](/text2vec/models/transformer.py#L8) + - [text2vec.models.TransformerDecoder](/text2vec/models/transformer.py#L78) + - [text2vec.models.RecurrentEncoder](/text2vec/models/sequential.py#L9) + - [text2vec.models.RecurrentDecoder](/text2vec/models/sequential.py#L65) #### Input and Word-Embeddings Components - - [text2vec.models.Tokenizer](/text2vec/models/components/feeder.py#L4) - - [text2vec.models.Embed](/text2vec/models/components/text_inputs.py#L4) - - [text2vec.models.TokenEmbed](/text2vec/models/components/text_inputs.py#L82) - - [text2vec.models.TextInput](/text2vec/models/components/feeder.py#L35) (DEPRECATED) + - [text2vec.models.Tokenizer](/text2vec/models/components/text_inputs.py#L5) + - [text2vec.models.Embed](/text2vec/models/components/text_inputs.py#L36) + - [text2vec.models.TokenEmbed](/text2vec/models/components/text_inputs.py#L116) #### Attention Components - - [text2vec.models.components.attention.ScaledDotAttention](/text2vec/models/components/attention.py#L4) - - [text2vec.models.components.attention.SingleHeadAttention](/text2vec/models/components/attention.py#L111) - - [text2vec.models.MultiHeadAttention](/text2vec/models/components/attention.py#L175) - - [text2vec.models.BahdanauAttention](/text2vec/models/components/attention.py#L53) + - [text2vec.models.components.attention.ScaledDotAttention](/text2vec/models/components/attention.py#L7) + - [text2vec.models.components.attention.SingleHeadAttention](/text2vec/models/components/attention.py#L115) + - [text2vec.models.MultiHeadAttention](/text2vec/models/components/attention.py#L179) + - [text2vec.models.BahdanauAttention](/text2vec/models/components/attention.py#L57) #### LSTM Components - - [text2vec.models.BidirectionalLSTM](/text2vec/models/components/recurrent.py#L4) + - [text2vec.models.BidirectionalLSTM](/text2vec/models/components/recurrent.py#L5) #### Pointwise Feedforward Components @@ -183,9 +95,10 @@ Text2vec includes a Python API with convenient classes for handling attention an #### General Layer Components - - [text2vec.models.components.utils.LayerNorm](/text2vec/models/components/utils.py#L5) + - [text2vec.models.components.utils.LayerNorm](/text2vec/models/components/utils.py#L6) - [text2vec.models.components.utils.TensorProjection](/text2vec/models/components/utils.py#L43) - - [text2vec.models.components.utils.PositionalEncder](/text2vec/models/components/utils.py#L76) + - [text2vec.models.components.utils.PositionalEncder](/text2vec/models/components/utils.py#L77) + - [text2vec.models.components.utils.VariationPositionalEncoder](/text2vec/models/components/utils.py#L122) #### Dataset Pre-processing diff --git a/setup.py b/setup.py index b8ad7ee..117b590 100644 --- a/setup.py +++ b/setup.py @@ -24,9 +24,4 @@ ] ), packages=find_packages(exclude=["bin"]), - entry_points={ - "console_scripts": [ - "text2vec_main=text2vec.bin.main:main", - ], - } ) diff --git a/text2vec/autoencoders.py b/text2vec/autoencoders.py index f520bdd..3c46748 100644 --- a/text2vec/autoencoders.py +++ b/text2vec/autoencoders.py @@ -4,13 +4,9 @@ import tensorflow as tf from tensorflow.keras import layers, Model -from text2vec.models.components.feeder import Tokenizer -from text2vec.models.components.text_inputs import TokenEmbed -from text2vec.models.components.text_inputs import Embed -from text2vec.models.transformer import TransformerEncoder -from text2vec.models.transformer import TransformerDecoder -from text2vec.models.sequential import RecurrentEncoder -from text2vec.models.sequential import RecurrentDecoder +from text2vec.models.components.text_inputs import TokenEmbed, Embed, Tokenizer +from text2vec.models.transformer import TransformerEncoder, TransformerDecoder +from text2vec.models.sequential import RecurrentEncoder, RecurrentDecoder class TransformerAutoEncoder(Model): @@ -124,7 +120,6 @@ def train_step(self, data): # pylint: disable=missing-function-docstring labels=targets.to_tensor(default_value=0) ) loss = loss * mask_decode - # loss = tf.math.reduce_sum(loss, axis=1) loss = tf.reduce_mean(loss) gradients = tape.gradient(loss, self.trainable_variables) @@ -147,7 +142,6 @@ def test_step(self, data): # pylint: disable=missing-function-docstring labels=targets.to_tensor(default_value=0) ) loss = loss * mask_decode - # loss = tf.math.reduce_sum(loss, axis=1) loss = tf.reduce_mean(loss) return {"loss": loss, **{m.name: m.result() for m in self.metrics}} @@ -272,11 +266,6 @@ def __init__(self, max_sequence_len: int, embedding_size: int, num_hidden: int = self.decode_layer = RecurrentDecoder(num_hidden=num_hidden, **params) def call(self, inputs, training: bool = False): # pylint: disable=missing-function-docstring - # tokens = self.tokenizer(tokens) - # x_enc, enc_mask, _ = self.embed_layer(tokens, training=training) - # x_enc, context, *states = self.encode_layer(x_enc, mask=enc_mask, training=training) - # return x_enc, context, enc_mask, states - encoding_text = inputs[0] decoding_text = inputs[1] if len(inputs) > 1 else encoding_text @@ -312,7 +301,6 @@ def train_step(self, data): # pylint: disable=missing-function-docstring labels=targets.to_tensor(default_value=0) ) loss = loss * mask_decode - # loss = tf.math.reduce_sum(loss, axis=1) loss = tf.reduce_mean(loss) gradients = tape.gradient(loss, self.trainable_variables) @@ -335,7 +323,6 @@ def test_step(self, data): # pylint: disable=missing-function-docstring labels=targets.to_tensor(default_value=0) ) loss = loss * mask_decode - # loss = tf.math.reduce_sum(loss, axis=1) loss = tf.reduce_mean(loss) return {"loss": loss, **{m.name: m.result() for m in self.metrics}} diff --git a/text2vec/models/__init__.py b/text2vec/models/__init__.py index c2c562f..74ed7f8 100644 --- a/text2vec/models/__init__.py +++ b/text2vec/models/__init__.py @@ -1,8 +1,7 @@ from text2vec.models.transformer import TransformerEncoder, TransformerDecoder from text2vec.models.sequential import RecurrentEncoder, RecurrentDecoder -from text2vec.models.components.feeder import TextInput, Tokenizer -from text2vec.models.components.text_inputs import Embed, TokenEmbed +from text2vec.models.components.text_inputs import Embed, TokenEmbed, Tokenizer from text2vec.models.components.attention import BahdanauAttention, MultiHeadAttention from text2vec.models.components.feed_forward import PositionWiseFFN from text2vec.models.components.recurrent import BidirectionalLSTM diff --git a/text2vec/models/components/feed_forward.py b/text2vec/models/components/feed_forward.py index ec80097..a050d98 100644 --- a/text2vec/models/components/feed_forward.py +++ b/text2vec/models/components/feed_forward.py @@ -39,5 +39,4 @@ def __init__(self, emb_dims: int): self.conv_outer = layers.Conv1D(filters=emb_dims, kernel_size=1, padding='same', use_bias=False) def call(self, x): - with tf.name_scope("PositionWiseFFN"): - return self.conv_outer(self.conv_inner(x)) + return self.conv_outer(self.conv_inner(x)) diff --git a/text2vec/models/components/feeder.py b/text2vec/models/components/feeder.py deleted file mode 100644 index 5ad773e..0000000 --- a/text2vec/models/components/feeder.py +++ /dev/null @@ -1,115 +0,0 @@ -import tensorflow as tf -from tensorflow.keras import layers - - -class Tokenizer(layers.Layer): - """String-splitting layer. - - Parameters - ---------- - sep : str, optional - The token to split the incoming strings by, by default ' '. - - Examples - -------- - ```python - import tensorflow as tf - from text2vec.models import Tokenizer - - text = tf.constant([ - "Sample string.", - "This is a second example." - ]) - tokenizer = Tokenizer() - tokenizer(text) - ``` - """ - - def __init__(self, sep: str = ' '): - super().__init__(name="Tokenizer") - self.sep = sep - - def call(self, corpus): - return tf.strings.split(corpus, self.sep) - - -class TextInput(layers.Layer): - """This layer handles the primary text feature transformations and word-embeddings to be passed off - to the sequence-aware parts of the encoder/decoder pipeline. - - Texts come in already tokenized. The tokens are transformed to integer index values from a - `tf.lookup.StaticHashTable` lookup table. The tokens are used to lookup word-embeddings and a sequence - mask is computed. - - The inputs are `tf.RaggedTensor` types, and after word-embeddings the tensor is made dense by padding to the - longest sequence length in the batch. - - In certain cases, only the word-embedding output is necessary, in which case `output_embeddings` can be set `True` - in the `__call__` method. This by-passes the padding and sequence masking steps. - - Parameters - ---------- - token_hash : dict - Token -> integer vocabulary lookup. - embedding_size : int - Dimensionality of the word-embeddings. - max_sequence_len : int - Longest sequence seen at training time. This layer ensures that no input sequences exceed this length. - - Examples - -------- - ```python - import tensorflow as tf - from text2vec.models import TextInput - - lookup = {'string': 0, 'is': 1, 'example': 2} - inputer = TextInput(token_hash=lookup, embedding_size=16, max_sequence_len=10) - - text = tf.ragged.constant([ - ["Sample", "string", "."], - ["This", "is", "a", "second", "example", "."] - ]) - sequences, seq_mask, time_steps = inputer(text) - - # get word-embeddings only - word_embeddings = inputer(text, output_embeddings=True) - ``` - """ - - def __init__(self, token_hash: dict, embedding_size: int, max_sequence_len: int): - super().__init__() - assert isinstance(token_hash, dict) - - self.num_labels = tf.constant(len(token_hash) + 1) - self.table = tf.lookup.StaticHashTable( - tf.lookup.KeyValueTensorInitializer( - keys=list(token_hash.keys()), - values=list(token_hash.values()), - value_dtype=tf.int32 - ), - default_value=max(token_hash.values()) + 1 - ) - self.embeddings = tf.Variable( - tf.random.uniform([self.num_labels, embedding_size], -1.0, 1.0), - name='embeddings', - dtype=tf.float32, - trainable=True - ) - self.max_len = tf.constant(max_sequence_len) - self.slicer = layers.Lambda(lambda x: x[:, :max_sequence_len], name="sequence-slice") - - def call(self, tokens: tf.RaggedTensor, output_embeddings: bool = False): - hashed = tf.ragged.map_flat_values(self.table.lookup, tokens) - hashed = self.slicer(hashed) - - x = tf.ragged.map_flat_values(tf.nn.embedding_lookup, self.embeddings, hashed) - if output_embeddings: - return x - - x = x.to_tensor(0) - x = x * tf.math.sqrt(tf.cast(tf.shape(self.embeddings)[-1], tf.float32)) # sqrt(embedding_size) - - seq_lengths = hashed.row_lengths() - time_steps = tf.cast(tf.reduce_max(seq_lengths), tf.int32) - mask = tf.sequence_mask(lengths=seq_lengths, maxlen=time_steps, dtype=tf.float32) - return x, mask, time_steps diff --git a/text2vec/models/components/text_inputs.py b/text2vec/models/components/text_inputs.py index 22c189c..26a1beb 100644 --- a/text2vec/models/components/text_inputs.py +++ b/text2vec/models/components/text_inputs.py @@ -2,6 +2,37 @@ from tensorflow.keras import layers +class Tokenizer(layers.Layer): + """String-splitting layer. + + Parameters + ---------- + sep : str, optional + The token to split the incoming strings by, by default ' '. + + Examples + -------- + ```python + import tensorflow as tf + from text2vec.models import Tokenizer + + text = tf.constant([ + "Sample string.", + "This is a second example." + ]) + tokenizer = Tokenizer() + tokenizer(text) + ``` + """ + + def __init__(self, sep: str = ' '): + super().__init__(name="Tokenizer") + self.sep = sep + + def call(self, corpus): + return tf.strings.split(corpus, self.sep) + + class Embed(layers.Layer): """This layer handles the primary text feature transformations and word-embeddings to be passed off to the sequence-aware parts of the encoder/decoder pipeline. diff --git a/text2vec/models/components/utils.py b/text2vec/models/components/utils.py index c9a6424..77e5165 100644 --- a/text2vec/models/components/utils.py +++ b/text2vec/models/components/utils.py @@ -83,11 +83,11 @@ class PositionalEncoder(layers.Layer): -------- ```python import tensorflow as tf - from text2vec.models import TextInput + from text2vec.models import TokenEmbed from text2vec.models import utils lookup = {'string': 0, 'is': 1, 'example': 2} - inputer = TextInput(token_hash=lookup, embedding_size=16, max_sequence_len=10) + inputer = TokenEmbed(token_hash=lookup, embedding_size=16, max_sequence_len=10) encoder = utils.PositionalEncoder(emb_dims=16, max_sequence_len=10) text = tf.ragged.constant([ @@ -134,11 +134,11 @@ class VariationPositionalEncoder(layers.Layer): -------- ```python import tensorflow as tf - from text2vec.models import TextInput + from text2vec.models import TokenEmbed from text2vec.models import utils lookup = {'string': 0, 'is': 1, 'example': 2} - inputer = TextInput(token_hash=lookup, embedding_size=16, max_sequence_len=10) + inputer = TokenEmbed(token_hash=lookup, embedding_size=16, max_sequence_len=10) encoder = utils.VariationPositionalEncoder(emb_dims=16, max_sequence_len=10) text = tf.ragged.constant([ diff --git a/text2vec/models/sequential.py b/text2vec/models/sequential.py index 0e1a964..b20cfff 100644 --- a/text2vec/models/sequential.py +++ b/text2vec/models/sequential.py @@ -24,11 +24,11 @@ class RecurrentEncoder(layers.Layer): -------- ```python import tensorflow as tf - from text2vec.models import TextInputs + from text2vec.models import TokenEmbed from text2vec.models import RecurrentEncoder lookup = {'string': 0, 'is': 1, 'example': 2} - inputer = TextInput(token_hash=lookup, embedding_size=16, max_sequence_len=10) + inputer = TokenEmbed(token_hash=lookup, embedding_size=16, max_sequence_len=10) encoder = RecurrentEncoder(max_sequence_len=10, num_hidden=8, input_keep_prob=0.75) text = tf.ragged.constant([ diff --git a/text2vec/models/transformer.py b/text2vec/models/transformer.py index 6d060f6..8743eab 100644 --- a/text2vec/models/transformer.py +++ b/text2vec/models/transformer.py @@ -1,11 +1,8 @@ from tensorflow.keras import layers -from .components.attention import BahdanauAttention -from .components.attention import MultiHeadAttention +from .components.attention import BahdanauAttention, MultiHeadAttention from .components.feed_forward import PositionWiseFFN -from .components.utils import VariationPositionalEncoder -from .components.utils import LayerNorm -from .components.utils import TensorProjection +from .components.utils import VariationPositionalEncoder, LayerNorm, TensorProjection class TransformerEncoder(layers.Layer): @@ -30,11 +27,11 @@ class TransformerEncoder(layers.Layer): -------- ```python import tensorflow as tf - from text2vec.models import TextInputs + from text2vec.models import TokenEmbed from text2vec.models import TransformerEncoder lookup = {'string': 0, 'is': 1, 'example': 2} - inputer = TextInput(token_hash=lookup, embedding_size=16, max_sequence_len=10) + inputer = TokenEmbed(token_hash=lookup, embedding_size=16, max_sequence_len=10) encoder = TransformerEncoder(max_sequence_len=10, embedding_size=16, input_keep_prob=0.75) text = tf.ragged.constant([ From 8c42a37085576af1ff3235f3639861337a8e84db Mon Sep 17 00:00:00 2001 From: Dave Date: Thu, 7 Jul 2022 13:52:56 -0400 Subject: [PATCH 08/10] doc updates --- README.md | 31 ++-- text2vec/__init__.py | 4 +- text2vec/training_tools.py | 291 ------------------------------------- 3 files changed, 20 insertions(+), 306 deletions(-) delete mode 100644 text2vec/training_tools.py diff --git a/README.md b/README.md index e60666f..592eaf9 100644 --- a/README.md +++ b/README.md @@ -5,20 +5,16 @@ Models for contextual embedding of arbitrary texts. ## Setup --- -To get started, one should have a flavor of TensorFlow installed, with -version `>=2.4.1`. One can run +To get started, one should have a flavor of TensorFlow installed, with version `>=2.4.1`. One can run ```bash pip install tensorflow>=2.4.1 ``` -If one wishes to run the examples, some additional dependencies -from HuggingFace will need to be installed. The full installation -looks like +If one wishes to run the examples, some additional dependencies from HuggingFace will need to be installed. The full installation looks like ```bash pip install tensorflow>=2.4.1 tokenizers datasets ``` -To install the core components as an import-able Python library -simply run +To install the core components as an import-able Python library simply run ```bash pip install git+https://github.com/brainsqueeze/text2vec.git @@ -51,7 +47,21 @@ This is an adapted bi-directional LSTM encoder-decoder model with a self-attenti Both models are trained using Adam SGD with the learning-rate decay program in [[2](https://arxiv.org/abs/1706.03762)]. -The pre-built auto-encoder models inherit from [tf.keras.Model](https://www.tensorflow.org/api_docs/python/tf/keras/Model), and as such they can be trained using the [fit method](https://www.tensorflow.org/api_docs/python/tf/keras/Model#fit). An example of training on Wikitext data is available in the [examples folder](./examples/trainers/wiki_transformer.py). This uses HuggingFace [tokenizers](https://huggingface.co/docs/tokenizers/python/latest/) and [datasets](https://huggingface.co/docs/datasets/master/). +The pre-built auto-encoder models inherit from [tf.keras.Model](https://www.tensorflow.org/api_docs/python/tf/keras/Model), and as such they can be trained using the [fit method](https://www.tensorflow.org/api_docs/python/tf/keras/Model#fit). Training examples are available in the [examples folder](./examples/trainers). This uses HuggingFace [tokenizers](https://huggingface.co/docs/tokenizers/python/latest/) and [datasets](https://huggingface.co/docs/datasets/master/). + +If you wish to run the example training scripts then you will need to clone the repository +```bash +git clone https://github.com/brainsqueeze/text2vec.git +``` +and then run either +```bash +python -m examples.trainers.news_transformer +``` +for the attention-based transformer, or +```bash +python -m examples.trainers.news_lstm +``` +for the LSTM-based encoder. These examples use the [Multi-News](https://github.com/Alex-Fabbri/Multi-News) dataset via [HuggingFace](https://huggingface.co/datasets/multi_news). ## Python API @@ -120,10 +130,7 @@ Text2vec includes a Python API with convenient classes for handling attention an ## Inference Demo --- -Once a model is fully trained then a demo API can be run, along with a small -UI to interact with the REST API. This demo attempts to use the trained model -to condense long bodies of text into the most important sentences, using the -inferred embedded context vectors. +Once a model is fully trained then a demo API can be run, along with a small UI to interact with the REST API. This demo attempts to use the trained model to condense long bodies of text into the most important sentences, using the inferred embedded context vectors. To start the model server simply run ```bash diff --git a/text2vec/__init__.py b/text2vec/__init__.py index cf966ff..8a065e2 100644 --- a/text2vec/__init__.py +++ b/text2vec/__init__.py @@ -2,12 +2,10 @@ from . import autoencoders from . import preprocessing from . import optimizer_tools -from . import training_tools __all__ = [ 'models', 'autoencoders', 'preprocessing', - 'optimizer_tools', - 'training_tools' + 'optimizer_tools' ] diff --git a/text2vec/training_tools.py b/text2vec/training_tools.py deleted file mode 100644 index 795abde..0000000 --- a/text2vec/training_tools.py +++ /dev/null @@ -1,291 +0,0 @@ -from typing import Dict, Union -import tensorflow as tf - -from text2vec.models import TextInput -from text2vec.models import Tokenizer -from text2vec.models import TokenEmbed -from text2vec.models import Embed -from text2vec.models import TransformerEncoder -from text2vec.models import TransformerDecoder -from text2vec.models import RecurrentEncoder -from text2vec.models import RecurrentDecoder - - -class EncodingModel(tf.keras.Model): - """Wrapper model class to combine the encoder-decoder training pipeline. - - Parameters - ---------- - token_hash : dict - Token -> integer vocabulary lookup. - max_sequence_len : int - Longest sequence seen at training time. - n_stacks : int, optional - Number of encoding blocks to chain, by default 1 - layers : int, optional - Number of layers in the multi-head-attention layer, by default 8 - num_hidden : int, optional - Dimensionality of hidden LSTM layer weights, by default 64 - input_keep_prob : float, optional - Value between 0 and 1.0 which determines `1 - dropout_rate`, by default 1.0 - hidden_keep_prob : float, optional - Hidden states dropout. Value between 0 and 1.0 which determines `1 - dropout_rate`, by default 1.0 - embedding_size : int, optional - Dimensionality of the word-embeddings, by default 64 - recurrent : bool, optional - Set to True to use the LSTM based model, otherwise defaults to attention based model, by default False - sep : str, optional - Token separator, by default ' ' - - Examples - -------- - ```python - import tensorflow as tf - from text2vec.training_tools import EncodingModel - - lookup = {'string': 0, 'is': 1, 'example': 2} - params = dict( - max_sequence_len=10, - embedding_size=16, - input_keep_prob=0.9, - hidden_keep_prob=0.75 - ) - model = EncodingModel(token_hash=lookup, layers=8, **params) - - text = tf.constant([ - "sample string .", - "this is a second example ." - ]) - y_hat, time_steps, targets, context_vectors = model(text, training=True, return_vectors=True) - ``` - """ - - def __init__(self, token_hash, max_sequence_len, n_stacks=1, layers=8, num_hidden=64, - input_keep_prob=1.0, hidden_keep_prob=1.0, embedding_size=64, recurrent=False, sep=' '): - super().__init__() - - params = dict( - max_sequence_len=max_sequence_len, - embedding_size=embedding_size, - input_keep_prob=input_keep_prob, - hidden_keep_prob=hidden_keep_prob - ) - self.embed_layer = TextInput( - token_hash=token_hash, - embedding_size=embedding_size, - max_sequence_len=max_sequence_len - ) - self.tokenizer = Tokenizer(sep) - - if recurrent: - self.encode_layer = RecurrentEncoder(num_hidden=num_hidden, **params) - self.decode_layer = RecurrentDecoder(num_hidden=num_hidden, **params) - else: - self.encode_layer = TransformerEncoder(n_stacks=n_stacks, layers=layers, **params) - self.decode_layer = TransformerDecoder(n_stacks=n_stacks, layers=layers, **params) - - def call(self, sentences, training=False, return_vectors=False): - tokens = self.tokenizer(sentences) # turn sentences into ragged tensors of tokens - - # turn incoming sentences into relevant tensor batches - with tf.name_scope('Encoding'): - x_enc, enc_mask, _ = self.embed_layer(tokens) - if not training: - return self.encode_layer(x_enc, mask=enc_mask, training=False) - x_enc, context, *states = self.encode_layer(x_enc, mask=enc_mask, training=True) - - with tf.name_scope('Decoding'): - batch_size = tokens.nrows() - - with tf.name_scope('targets'): - eos = tf.fill([batch_size], value='', name='eos-tag') - eos = tf.expand_dims(eos, axis=-1, name='eos-tag-expand') - - targets = tf.concat([tokens, eos], 1, name='eos-concat') - targets = tf.ragged.map_flat_values(self.embed_layer.table.lookup, targets) - targets = self.embed_layer.slicer(targets) - - with tf.name_scope('decode-tokens'): - bos = tf.fill([batch_size], value='', name='bos-tag') - bos = tf.expand_dims(bos, axis=-1, name='bos-tag-expand') - - dec_tokens = tf.concat([bos, tokens], -1, name='bos-concat') - x_dec, dec_mask, dec_time_steps = self.embed_layer(dec_tokens) - x_out = self.decode_layer( - x_enc=x_enc, - enc_mask=enc_mask, - x_dec=x_dec, - dec_mask=dec_mask, - context=context, - attention=self.encode_layer.attention, - training=training, - initial_state=states - ) - x_out = tf.tensordot(x_out, self.embed_layer.embeddings, axes=[2, 1]) - - if return_vectors: - return x_out, dec_time_steps, targets.to_tensor(default_value=0), context - return x_out, dec_time_steps, targets.to_tensor(default_value=0) - - @tf.function(input_signature=[tf.TensorSpec(shape=[None], dtype=tf.string)]) - def embed(self, sentences): - """Takes batches of free text and returns context vectors for each example. - - Parameters - ---------- - sentences : tf.Tensor - Tensor of dtype tf.string. - - Returns - ------- - tf.Tensor - Context vectors of shape (batch_size, embedding_size) - """ - - tokens = self.tokenizer(sentences) # turn sentences into ragged tensors of tokens - x_enc, enc_mask, _ = self.embed_layer(tokens) - return self.encode_layer(x_enc, mask=enc_mask, training=False) - - @tf.function(input_signature=[tf.TensorSpec(shape=[None], dtype=tf.string)]) - def token_embed(self, sentences): - """Takes batches of free text and returns word embeddings along with the associate token. - - Parameters - ---------- - sentences : tf.Tensor - Tensor of dtype tf.string. - - Returns - ------- - (tf.Tensor, tf.Tensor) - Tuple of (tokens, word_embeddings) with shapes (batch_size, max_sequence_len) - and (batch_size, max_sequence_len, embedding_size) respectively. - """ - - tokens = self.tokenizer(sentences) # turn sentences into ragged tensors of tokens - return tokens.to_tensor(''), self.embed_layer(tokens, output_embeddings=True).to_tensor(0) - - -class ServingModel(tf.keras.Model): - """Wrapper class for packaging final layers prior to saving. - - Parameters - ---------- - embed_layer : Union[TokenEmbed, Embed] - text2vec `TokenEmbed` or `Embed` layer - encode_layer : Union[TransformerEncoder, RecurrentEncoder] - text2vec `TransformerEncoder` or `RecurrentEncoder` layer - tokenizer : Tokenizer - text2vec `Tokenizer` layer - """ - - def __init__(self, embed_layer: Union[TokenEmbed, Embed], - encode_layer: Union[TransformerEncoder, RecurrentEncoder], tokenizer: Tokenizer): - super().__init__() - self.embed_layer = embed_layer - self.tokenizer = tokenizer - self.encode_layer = encode_layer - - @tf.function(input_signature=[tf.TensorSpec(shape=[None], dtype=tf.string)]) - def embed(self, sentences) -> Dict[str, tf.Tensor]: - """Takes batches of free text and returns context vectors for each example. - - Parameters - ---------- - sentences : tf.Tensor - Tensor of dtype tf.string. - - Returns - ------- - Dict[str, tf.Tensor] - Attention vector and hidden state sequences with shapes (batch_size, embedding_size) - and (batch_size, max_sequence_len, embedding_size) respectively. - """ - - tokens = self.tokenizer(sentences) # turn sentences into ragged tensors of tokens - x_enc, enc_mask, _ = self.embed_layer(tokens) - sequences, attention = self.encode_layer(x_enc, mask=enc_mask, training=False) - return {"sequences": sequences, "attention": attention} - - @tf.function(input_signature=[tf.TensorSpec(shape=[None], dtype=tf.string)]) - def token_embed(self, sentences) -> Dict[str, tf.Tensor]: - """Takes batches of free text and returns word embeddings along with the associate token. - - Parameters - ---------- - sentences : tf.Tensor - Tensor of dtype tf.string. - - Returns - ------- - Dict[str, tf.Tensor] - Padded tokens and embedding vectors with shapes (batch_size, max_sequence_len) - and (batch_size, max_sequence_len, embedding_size) respectively. - """ - - tokens = self.tokenizer(sentences) # turn sentences into ragged tensors of tokens - return { - "tokens": tokens.to_tensor(''), - "embeddings": self.embed_layer.get_embedding(tokens).to_tensor(0) - } - - -def sequence_cost(target_sequences, sequence_logits, num_labels, smoothing=False): - """Sequence-to-sequence cost function with optional label smoothing. - - Parameters - ---------- - target_sequences : tf.Tensor - Expected token sequences as lookup IDs (batch_size, max_sequence_len) - sequence_logits : [type] - Computed logits for predicted tokens (batch_size, max_sequence_len, embedding_size) - num_labels : int - Vocabulary look up size - smoothing : bool, optional - Set to True to smooth labels, this increases regularization while increasing training time, by default False - - Returns - ------- - tf.float32 - Loss value averaged over examples. - """ - - with tf.name_scope('Cost'): - if smoothing: - smoothing = 0.1 - targets = tf.one_hot(target_sequences, depth=num_labels, on_value=1.0, off_value=0.0, axis=-1) - loss = tf.keras.losses.binary_crossentropy( - y_true=targets, - y_pred=sequence_logits, - from_logits=True, - label_smoothing=smoothing - ) - else: - loss = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=sequence_logits, labels=target_sequences) - - loss = tf.reduce_mean(loss) - return loss - - -def vector_cost(context_vectors): - """Cost constraint on the cosine similarity of context vectors. Diagonal elements (self-context) - are coerced to be closer to 1 (self-consistency). Off-diagonal elements are pushed toward 0, - indicating not contextually similar. - - Parameters - ---------- - context_vectors : tf.Tensor - (batch_size, embedding_size) - - Returns - ------- - tf.float32 - cosine similarity constraint loss - """ - - with tf.name_scope('VectorCost'): - rows = tf.shape(context_vectors)[0] - context_vectors = tf.linalg.l2_normalize(context_vectors, axis=-1) - cosine = tf.tensordot(context_vectors, tf.transpose(context_vectors), axes=[1, 0]) - identity = tf.eye(rows) - return tf.reduce_mean((identity - cosine) ** 2) From a18f607abd769a5b5c44d373eaafb117594802f1 Mon Sep 17 00:00:00 2001 From: Dave Date: Thu, 7 Jul 2022 15:19:34 -0400 Subject: [PATCH 09/10] Update demo app --- demo/api.py | 127 +++++++++++++++++++++++++ demo/{context.html => html/index.html} | 6 +- demo/js/condense.js | 15 +-- demo/package.json | 19 ---- 4 files changed, 139 insertions(+), 28 deletions(-) create mode 100644 demo/api.py rename demo/{context.html => html/index.html} (78%) delete mode 100644 demo/package.json diff --git a/demo/api.py b/demo/api.py new file mode 100644 index 0000000..8bbd9c2 --- /dev/null +++ b/demo/api.py @@ -0,0 +1,127 @@ +from typing import List, Union +from math import pi +import argparse +import json +import re + +from flask import Flask, request, Response, send_from_directory +from flask_cors import cross_origin +from tornado.log import enable_pretty_logging +from tornado.httpserver import HTTPServer +from tornado.wsgi import WSGIContainer +from tornado.ioloop import IOLoop +import tornado.autoreload +# from tornado import web +import tornado + +import tensorflow as tf +from tensorflow.keras import models, Model +from tokenizers import Tokenizer + +app = Flask(__name__, static_url_path="", static_folder="./") +parser = argparse.ArgumentParser() +parser.add_argument("--model_dir", type=str, help="Directory containing serialized model and tokenizer", required=True) +args = parser.parse_args() + +model: Model = models.load_model(f"{args.model_dir}/saved_model") +tokenizer: Tokenizer = Tokenizer.from_file(f"{args.model_dir}/tokenizer.json") + + +def responder(results, error, message): + """Boilerplate Flask response item. + + Parameters + ---------- + results : dict + API response + error : int + Error code + message : str + Message to send to the client + + Returns + ------- + flask.Reponse + """ + + assert isinstance(results, dict) + results["message"] = message + results = json.dumps(results, indent=2) + + return Response( + response=results, + status=error, + mimetype="application/json" + ) + + +def tokenize(text: Union[str, List[str]]) -> List[str]: + if isinstance(text, str): + return [' '.join(tokenizer.encode(text).tokens)] + return [' '.join(batch.tokens) for batch in tokenizer.encode_batch(text)] + + +def get_summaries(paragraphs: List[str]): + context = tf.concat([ + model.embed(batch)["attention"] + for batch in tf.data.Dataset.from_tensor_slices(paragraphs).batch(32) + ], axis=0) + doc_vector = model.embed(tf.strings.reduce_join(paragraphs, separator=' ', keepdims=True))["attention"] + cosine = tf.tensordot(tf.math.l2_normalize(context, axis=1), tf.math.l2_normalize(doc_vector, axis=1), axes=[-1, 1]) + cosine = tf.clip_by_value(cosine, -1, 1) + likelihoods = tf.nn.softmax(180 - tf.math.acos(cosine) * (180 / pi), axis=0) + return likelihoods + + +@app.route("/") +def root(): + return send_from_directory(directory="./html/", path="index.html") + + +@app.route("/summarize", methods=["GET", "POST"]) +# @cross_origin(origins=['*'], allow_headers=['Content-Type', 'Authorization']) +def summarize(): + if request.is_json: + payload = request.json + else: + payload = request.values + + text = payload.get("text", "") + if not text: + return responder(results={}, error=400, message="No text provided") + + paragraphs = [p for p in re.split(r"\n{1,}", text) if p.strip()] + if len(paragraphs) < 2: + return responder(results={"text": paragraphs}, error=400, message="Insufficient amount of text provided") + + tokenized = tokenize(paragraphs) + likelihoods = get_summaries(tokenized) + likelihoods = tf.squeeze(likelihoods) + cond = tf.where(likelihoods > tf.math.reduce_mean(likelihoods) + tf.math.reduce_std(likelihoods)).numpy().flatten() + output = [{ + "text": paragraphs[idx], + "score": float(likelihoods[idx]) + } for idx in cond] + + results = {"data": output} + return responder(results=results, error=200, message="Success") + + +def serve(port: int = 9090, debug: bool = False): + http_server = HTTPServer(WSGIContainer(app)) + http_server.listen(port) + enable_pretty_logging() + + io_loop = IOLoop.current() + if debug: + tornado.autoreload.start(check_time=500) + print("Listening to port", port, flush=True) + + try: + io_loop.start() + except KeyboardInterrupt: + pass + + +if __name__ == '__main__': + serve() diff --git a/demo/context.html b/demo/html/index.html similarity index 78% rename from demo/context.html rename to demo/html/index.html index 66b26f7..0c9353f 100644 --- a/demo/context.html +++ b/demo/html/index.html @@ -1,14 +1,14 @@ - +
-

Contextual condenser demo

+

Summary extraction demo

@@ -17,7 +17,7 @@

Contextual condenser demo

- + diff --git a/demo/js/condense.js b/demo/js/condense.js index 90f56b1..9f39b6d 100644 --- a/demo/js/condense.js +++ b/demo/js/condense.js @@ -10,16 +10,19 @@ function formatOutput (ajaxData, container) { for (var i = 0; i < data.length; i++) { let text = data[i].text - let score = data[i].relevanceScore - let lightness = (1 - score) * 100.0 + let score = data[i].score + let lightness = score * 100. - outputHtml += `

${text}

` + if (score > 0.5) + outputHtml += `

${text}

` + else + outputHtml += `

${text}

` } container.append(outputHtml) } } -$(document).ready(function () { +$(function () { $('#clear').on('click', function () { $('.text-input.main').val('') $('.response-container').empty() @@ -27,12 +30,12 @@ $(document).ready(function () { $('#go').on('click', function () { let text = $('.text-input.main').val() - let data = { body: text } + let data = { text: text } $('.response-container').empty() $.ajax({ - url: 'http://localhost:8008/condense', + url: '/summarize', data: data, type: 'POST', dataType: 'json', diff --git a/demo/package.json b/demo/package.json deleted file mode 100644 index 077b0ac..0000000 --- a/demo/package.json +++ /dev/null @@ -1,19 +0,0 @@ -{ - "name": "text2vec-demo", - "version": "1.0.0", - "description": "Demonstrates the ability of text2vec to condense large bodies of text using the contextual embeddings.", - "main": "index.js", - "scripts": { - "test": "echo \"Error: no test specified\" && exit 1" - }, - "author": "Dave Hollander", - "license": "BSD-2-Clause", - "devDependencies": { - "eslint": "^5.7.0", - "eslint-config-standard": "^12.0.0", - "eslint-plugin-import": "^2.14.0", - "eslint-plugin-node": "^7.0.1", - "eslint-plugin-promise": "^4.0.1", - "eslint-plugin-standard": "^4.0.0" - } -} From 3b974a5a82d03376999a5e4b8f70997b404f900f Mon Sep 17 00:00:00 2001 From: Dave Date: Thu, 7 Jul 2022 15:33:42 -0400 Subject: [PATCH 10/10] updated demo, remove cruft --- README.md | 13 +- demo/api.py | 3 - examples/trainers/news_transformer.py | 4 +- setup.py | 2 - text2vec/bin/__init__.py | 0 text2vec/bin/main.py | 266 -------------------------- text2vec/bin/serving_tools.py | 53 ----- text2vec/bin/text_summarize.py | 232 ---------------------- text2vec/bin/utils.py | 26 --- 9 files changed, 10 insertions(+), 589 deletions(-) delete mode 100644 text2vec/bin/__init__.py delete mode 100644 text2vec/bin/main.py delete mode 100644 text2vec/bin/serving_tools.py delete mode 100644 text2vec/bin/text_summarize.py delete mode 100644 text2vec/bin/utils.py diff --git a/README.md b/README.md index 592eaf9..1b808ae 100644 --- a/README.md +++ b/README.md @@ -130,14 +130,17 @@ Text2vec includes a Python API with convenient classes for handling attention an ## Inference Demo --- -Once a model is fully trained then a demo API can be run, along with a small UI to interact with the REST API. This demo attempts to use the trained model to condense long bodies of text into the most important sentences, using the inferred embedded context vectors. - +Trained text2vec models can be demonstrated from a lightweight app included in this repository. The demo runs extractive summarization from long bodies of text using the attention vectors of the encoding latent space. To get started, you will need to clone the repository and then install additional dependencies: +```bash +git clone https://github.com/brainsqueeze/text2vec.git +cd text2vec +pip install flask tornado +``` To start the model server simply run ```bash -text2vec_main --run=infer --yaml_config=/path/to/config.yml +python demo/api.py --model_dir /absolute/saved_model/parent/dir ``` -A demonstration webpage is included in [demo](demo) at -[context.html](demo/context.html). +The `model_dir` CLI parameter must be an absolute path to the directory containing the `/saved_model` folder and the `tokenizer.json` file from a text2vec model with an `embed` signature. A demonstration app is served on port 9090. ## References --- diff --git a/demo/api.py b/demo/api.py index 8bbd9c2..e238f6b 100644 --- a/demo/api.py +++ b/demo/api.py @@ -5,13 +5,11 @@ import re from flask import Flask, request, Response, send_from_directory -from flask_cors import cross_origin from tornado.log import enable_pretty_logging from tornado.httpserver import HTTPServer from tornado.wsgi import WSGIContainer from tornado.ioloop import IOLoop import tornado.autoreload -# from tornado import web import tornado import tensorflow as tf @@ -79,7 +77,6 @@ def root(): @app.route("/summarize", methods=["GET", "POST"]) -# @cross_origin(origins=['*'], allow_headers=['Content-Type', 'Authorization']) def summarize(): if request.is_json: payload = request.json diff --git a/examples/trainers/news_transformer.py b/examples/trainers/news_transformer.py index 5d8777b..95941f2 100644 --- a/examples/trainers/news_transformer.py +++ b/examples/trainers/news_transformer.py @@ -37,7 +37,7 @@ def train_tokenizer() -> Tuple[tokenizers.Tokenizer, tf.data.Dataset]: pre_tokenizers.Digits(individual_digits=False) ]) - dataset = datasets.load_dataset("multi_news", split="train") + dataset = datasets.load_dataset("multi_news", split="test") def batch_iterator(batch_size=1000): for i in range(0, len(dataset), batch_size): @@ -126,7 +126,7 @@ def main(save_path: str): ) ) ], - epochs=10 + epochs=2 ) model.save( diff --git a/setup.py b/setup.py index 117b590..7d85d2a 100644 --- a/setup.py +++ b/setup.py @@ -18,8 +18,6 @@ extras_require=dict( serving=[ "flask", - "flask-cors", - "nltk", "tornado" ] ), diff --git a/text2vec/bin/__init__.py b/text2vec/bin/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/text2vec/bin/main.py b/text2vec/bin/main.py deleted file mode 100644 index 1086eca..0000000 --- a/text2vec/bin/main.py +++ /dev/null @@ -1,266 +0,0 @@ -import itertools -import argparse -import os - -import yaml - -import numpy as np -import tensorflow as tf -from tensorboard.plugins import projector - -from text2vec.training_tools import EncodingModel -from text2vec.training_tools import ServingModel -from text2vec.training_tools import sequence_cost -from text2vec.training_tools import vector_cost -from text2vec.optimizer_tools import RampUpDecaySchedule -from text2vec.preprocessing.text import clean_and_split -from text2vec.preprocessing import utils as data_tools -from . import utils - -root = os.path.dirname(os.path.abspath(__file__)) -os.environ['KMP_DUPLICATE_LIB_OK'] = 'True' - - -def train(model_folder, num_tokens=10000, embedding_size=256, num_hidden=128, max_allowed_seq=-1, - layers=8, batch_size=32, num_epochs=10, data_files=None, model_path=".", use_attention=False, - eval_sentences=None, orthogonal_cost=False): - """Core training algorithm. - - Parameters - ---------- - model_folder : str - Name of the folder to create for the trained model - num_tokens : int, optional - Number of vocab tokens to keep from the training corpus, by default 10000 - embedding_size : int, optional - Size of the word-embedding dimensions, by default 256 - num_hidden : int, optional - Number of hidden model dimensions, by default 128 - max_allowed_seq : int, optional - The maximum sequence length allowed, model will truncate if longer, by default -1 - layers : int, optional - Number of multi-head attention mechanisms for transformer model, by default 8 - batch_size : int, optional - Size of each mini-batch, by default 32 - num_epochs : int, optional - Number of training epochs, by default 10 - data_files : list, optional - List of absolute paths to training data sets, by default None - model_path : str, optional - Valid path to where the model will be saved, by default "." - use_attention : bool, optional - Set to True to use the self-attention only model, by default False - eval_sentences : List, optional - List of sentences to check the context angles, by default None - orthogonal_cost : bool, optional - Set to True to add a cost to mutually parallel context vector, by default False - - Returns - ------- - str - Model checkpoint file path. - """ - - # GPU config - for gpu in tf.config.experimental.list_physical_devices('GPU'): - # tf.config.experimental.set_memory_growth(gpu, True) - tf.config.experimental.set_memory_growth(gpu, False) - tf.config.set_soft_device_placement(True) - - log_dir = f"{model_path}/{model_folder}" if model_path else f"{root}/../../text2vec/{model_folder}" - if not os.path.exists(log_dir): - os.mkdir(log_dir) - - utils.log("Fetching corpus and creating data pipeline") - corpus = data_tools.load_text_files(data_files=data_files, max_length=max_allowed_seq) - - utils.log("Fitting embedding lookup", end="...") - hash_map, max_seq_len, train_set_size = data_tools.get_top_tokens(corpus, n_top=num_tokens) - print(f"{train_set_size} sentences. max sequence length: {max_seq_len}") - - with open(log_dir + "/metadata.tsv", "w") as tsv: - for token, _ in sorted(hash_map.items(), key=lambda s: s[-1]): - # since tensorflow converts strings to byets we will decode from UTF-8 here for display purposes - tsv.write(f"{token.decode('utf8', 'replace')}\n") - tsv.write("\n") - - utils.log("Building computation graph") - log_step = (train_set_size // batch_size) // 25 - dims = embedding_size - - params = dict( - max_sequence_len=max_seq_len, - embedding_size=dims, - input_keep_prob=0.9, - hidden_keep_prob=0.75 - ) - if use_attention: - model = EncodingModel(token_hash=hash_map, layers=layers, **params) - else: - model = EncodingModel(token_hash=hash_map, num_hidden=num_hidden, recurrent=True, **params) - - warmup_steps = max(train_set_size // batch_size, 4000) - learning_rate = RampUpDecaySchedule(embedding_size=dims, warmup_steps=warmup_steps) - optimizer = tf.keras.optimizers.Adam(learning_rate=learning_rate) - train_loss = tf.keras.metrics.Mean('train-loss', dtype=tf.float32) - - def compute_loss(sentences): - y_hat, time_steps, targets, vectors = model(sentences, training=True, return_vectors=True) - loss_val = sequence_cost( - target_sequences=targets, - sequence_logits=y_hat[:, :time_steps], - num_labels=model.embed_layer.num_labels, - smoothing=False - ) - - if orthogonal_cost: - return loss_val + vector_cost(context_vectors=vectors) - return loss_val - - @tf.function(input_signature=[tf.TensorSpec(shape=(None,), dtype=tf.string)]) - def train_step(sentences): - loss_val = compute_loss(sentences) - gradients = tf.gradients(loss_val, model.trainable_variables) - optimizer.apply_gradients(zip(gradients, model.trainable_variables)) - train_loss(loss_val) # log the loss value to TensorBoard - - model_file_name = None - if isinstance(eval_sentences, list) and len(eval_sentences) > 1: - test_sentences = eval_sentences - else: - test_sentences = ["The movie was great!", "The movie was terrible."] - test_tokens = [' '.join(clean_and_split(text)) for text in test_sentences] - - summary_writer_train = tf.summary.create_file_writer(log_dir + "/training") - summary_writer_dev = tf.summary.create_file_writer(log_dir + "/validation") - checkpoint = tf.train.Checkpoint(EmbeddingModel=model, optimizer=optimizer) - checkpoint_manager = tf.train.CheckpointManager(checkpoint, log_dir, max_to_keep=5) - - # add word labels to the projector - config = projector.ProjectorConfig() - # pylint: disable=no-member - embeddings_config = config.embeddings.add() - - checkpoint_manager.save() - reader = tf.train.load_checkpoint(log_dir) - embeddings_config.tensor_name = [key for key in reader.get_variable_to_shape_map() if "embedding" in key][0] - embeddings_config.metadata_path = log_dir + "/metadata.tsv" - projector.visualize_embeddings(logdir=log_dir + "/training", config=config) - - step = 1 - for epoch in range(num_epochs): - try: - corpus = corpus.unbatch() - except ValueError: - print("Corpus not batched") - corpus = corpus.shuffle(train_set_size) - corpus = corpus.batch(batch_size).prefetch(10) # pre-fetch 10 batches for queuing - - print(f"\t Epoch: {epoch + 1}") - i = 1 - train_loss.reset_states() - - for x in corpus: - if step == 1: - tf.summary.trace_on(graph=True, profiler=False) - - train_step(x) - with summary_writer_train.as_default(): - if step == 1: - tf.summary.trace_export(name='graph', step=1, profiler_outdir=log_dir) - tf.summary.trace_off() - summary_writer_train.flush() - - if i % log_step == 0: - print(f"\t\t iteration {i} - loss: {train_loss.result()}") - tf.summary.scalar(name='loss', data=train_loss.result(), step=step) - tf.summary.scalar(name='learning-rate', data=learning_rate.callback(step=step), step=step) - summary_writer_train.flush() - train_loss.reset_states() - i += 1 - step += 1 - - vectors = model.embed(test_tokens) - angles = utils.compute_angles(vectors.numpy()) - - with summary_writer_dev.as_default(): - for idx, (i, j) in enumerate(itertools.combinations(range(len(test_sentences)), r=2), start=1): - angle = angles[i, j] - print(f"The angle between '{test_sentences[i]}' and '{test_sentences[j]}' is {angle} degrees") - - # log the angle to tensorboard - desc = f"'{test_sentences[i]}' : '{test_sentences[j]}'" - tf.summary.scalar(f'similarity-angle/{idx}', angle, step=step, description=desc) - summary_writer_dev.flush() - model_file_name = checkpoint_manager.save() - - utils.log("Saving a frozen model") - serve_model_ = ServingModel(embed_layer=model.embed_layer, encode_layer=model.encode_layer, sep=' ') - tf.saved_model.save( - obj=serve_model_, - export_dir=f"{log_dir}/frozen/1", - signatures={"serving_default": serve_model_.embed, "token_embed": serve_model_.token_embed} - ) - - utils.log("Reloading frozen model and comparing output to in-memory model") - test = tf.saved_model.load(f"{log_dir}/frozen/1") - test_model = test.signatures["serving_default"] - test_output = test_model(tf.constant(test_tokens))["output_0"] - utils.log(f"Outputs on CV set are approximately the same?: {np.allclose(test_output, model.embed(test_tokens))}") - return model_file_name - - -def main(): - """Training and inferencing entrypoint for CLI. - - Raises - ------ - NotImplementedError - Raised if a `run` mode other than `train` or `infer` are passed. - """ - - parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) - parser.add_argument("--run", choices=["train", "infer"], help="Run type.", required=True) - parser.add_argument("--attention", action='store_true', help="Set to use attention transformer model.") - parser.add_argument("--orthogonal", action='store_true', help="Set to add a cost to mutually parallel contexts.") - parser.add_argument("--yaml_config", type=str, help="Path to a valid training config YAML file.", required=True) - args = parser.parse_args() - - config_path = args.yaml_config - if config_path.startswith("${HOME}"): - config_path = config_path.replace('${HOME}', os.getenv('HOME')) - elif config_path.startswith("$HOME"): - config_path = config_path.replace('$HOME', os.getenv('HOME')) - - config = yaml.safe_load(open(config_path, 'r')) - training_config = config.get("training", {}) - model_config = config.get("model", {}) - model_params = model_config.get("parameters", {}) - - if args.run == "train": - train( - model_folder=model_config["name"], - use_attention=args.attention, - num_tokens=training_config.get("tokens", 10000), - max_allowed_seq=training_config.get("max_sequence_length", 512), - embedding_size=model_params.get("embedding", 128), - num_hidden=model_params.get("hidden", 128), - layers=model_params.get("layers", 8), - batch_size=training_config.get("batch_size", 32), - num_epochs=training_config.get("epochs", 20), - data_files=training_config.get("data_files"), - model_path=model_config.get("storage_dir", "."), - eval_sentences=training_config.get("eval_sentences"), - orthogonal_cost=args.orthogonal - ) - elif args.run == "infer": - os.environ["MODEL_PATH"] = f'{model_config.get("storage_dir", ".")}/{model_config["name"]}' - from .text_summarize import run_server - run_server(port=8008) - else: - raise NotImplementedError("Only training and inferencing is enabled right now.") - return - - -if __name__ == '__main__': - main() diff --git a/text2vec/bin/serving_tools.py b/text2vec/bin/serving_tools.py deleted file mode 100644 index 8305c13..0000000 --- a/text2vec/bin/serving_tools.py +++ /dev/null @@ -1,53 +0,0 @@ -import os - -from nltk.tokenize import sent_tokenize - -import tensorflow as tf -from text2vec.preprocessing.text import normalize_text, clean_and_split - - -class Embedder(): - """Wrapper class which handles contextual embedding of documents. - """ - - def __init__(self): - log_dir = f"{os.environ['MODEL_PATH']}/frozen/1" - self.__model = tf.saved_model.load(log_dir) - - @staticmethod - def __get_sentences(text): - data = [(sent, ' '.join(clean_and_split(normalize_text(sent)))) for sent in sent_tokenize(text)] - data = [(orig, clean) for orig, clean in data if len(clean.split()) >= 5] - original, clean = map(list, zip(*data)) - return original, clean - - def __normalize(self, vectors: tf.Tensor): - return tf.math.l2_normalize(vectors, axis=-1).numpy() - - def __doc_vector(self, doc: tf.Tensor): - net_vector = tf.reduce_sum(doc, axis=0) - return self.__normalize(net_vector) - - def __embed(self, corpus: list): - return self.__model.embed(corpus) - - def embed(self, text: str): - """String preparation and embedding. Returns the context vector representing the input document. - - Parameters - ---------- - text : str - - Returns - ------- - (list, tf.Tensor, tf.Tensor) - ( - Segmented sentences, - L2-normalized context vectors (num_sentences, embedding_size), - Single unit vector representing the entire document (embedding_size,) - ) - """ - - sentences, clean_sentences = self.__get_sentences(text) - vectors = self.__embed(clean_sentences) - return sentences, self.__normalize(vectors), self.__doc_vector(vectors) diff --git a/text2vec/bin/text_summarize.py b/text2vec/bin/text_summarize.py deleted file mode 100644 index b9a6428..0000000 --- a/text2vec/bin/text_summarize.py +++ /dev/null @@ -1,232 +0,0 @@ -import time -import json - -from flask import Flask, request, Response -from flask_cors import cross_origin - -from tornado.httpserver import HTTPServer -from tornado.wsgi import WSGIContainer -from tornado.ioloop import IOLoop -import tornado.autoreload -import tornado - -import numpy as np -from .serving_tools import Embedder - -app = Flask(__name__) -model = Embedder() - - -def responder(results, error, message): - """Boilerplate Flask response item. - - Parameters - ---------- - results : dict - API response - error : int - Error code - message : str - Message to send to the client - - Returns - ------- - flask.Reponse - """ - - assert isinstance(results, dict) - results["message"] = message - results = json.dumps(results, indent=2) - - return Response( - response=results, - status=error, - mimetype="application/json" - ) - - -def cosine_similarity_sort(net_vector, embedding_matrix): - """ - Computes the cosine similarity scores and then returns - the sorted results - - Parameters - ---------- - net_vector : np.ndarray - The context vector for the entire document - embedding_matrix : np.ndarray - The context vectors (row vectors) for each constituent body of text - - Returns - ------- - (ndarray, ndarray) - (sorted order of documents, cosine similarity scores) - """ - - similarity = np.dot(embedding_matrix, net_vector) - similarity = np.clip(similarity, -1, 1) - # sort = np.argsort(1 - similarity) - sort = np.argsort(similarity - 1) - - return sort, similarity.flatten()[sort] - - -def angle_from_cosine(cosine_similarity): - """ - Computes the angles in degrees from cosine similarity scores - - Parameters - ---------- - cosine_similarity : np.ndarray - - Returns - ------- - ndarray - Cosine angles (num_sentences,) - """ - - return np.arccos(cosine_similarity) * (180 / np.pi) - - -def choose(sentences, scores, embeddings): - """ - Selects the best constituent texts from the similarity scores - - Parameters - ---------- - sentences : np.ndarray - Array of the input texts, sorted by scores. - scores : np.ndarray - Cosine similarity scores, sorted - embeddings : np.ndarray - Embedding matrix for input texts, sorted by scores - - Returns - ------- - (np.ndarray, np.ndarray, np.ndarray) - (best sentences sorted, best scores sorted, best embeddings sorted) - """ - - if scores.shape[0] == 1: - return sentences, scores, embeddings - - angles = angle_from_cosine(scores) - cut = angles < angles.mean() - angles.std() - return sentences[cut], scores[cut], embeddings[cut] - - -def text_pass_filter(texts, texts_embeddings, net_vector): - """ - Runs the scoring + filtering process on input texts - - Parameters - ---------- - texts : np.ndarray - Input texts. - texts_embeddings : np.ndarray - Context embedding matrix for input texts. - net_vector : np.ndarray - The context vector for the entire document - - Returns - ------- - (np.ndarray, np.ndarray, np.ndarray) - (best sentences sorted, best scores sorted, best embeddings sorted) - """ - - sorted_order, scores = cosine_similarity_sort(net_vector=net_vector, embedding_matrix=texts_embeddings) - texts = np.array(texts)[sorted_order] - filtered_texts, filtered_scores, filtered_embeddings = choose( - sentences=texts, - scores=scores, - embeddings=texts_embeddings[sorted_order] - ) - - return filtered_texts, filtered_scores, filtered_embeddings - - -def softmax(logits): - """ - Computes the softmax of the input logits. - - Parameters - ---------- - logits : np.ndarray - - Returns - ------- - np.ndarray - Softmax output array with the same shape as the input. - """ - - soft = np.exp(logits) - soft[np.isinf(soft)] = 1e10 - soft /= np.sum(soft, axis=0) - soft = np.clip(soft, 0.0, 1.0) - return soft - - -@app.route('/condense', methods=['POST', 'GET']) -@cross_origin(origins=['*'], allow_headers=['Content-Type', 'Authorization']) -def compute(): - """ - Main Flask handler function - - Returns - ------- - flask.Response - """ - - j = request.get_json() - if j is None: - j = request.args - if not j: - j = request.form - - st = time.time() - body = j.get("body", "") - if not body: - results = { - "elapsed_time": time.time() - st, - "data": None - } - return responder(results=results, error=400, message="No text provided") - - # get the embedding vectors for each sentence in the document - sentences, vectors, doc_vector = model.embed(body) - top_sentences, top_scores, _ = text_pass_filter(texts=sentences, texts_embeddings=vectors, net_vector=doc_vector) - - results = { - "elapsed_time": time.time() - st, - "data": [{ - "text": text, - "relevanceScore": score - } for text, score in zip(top_sentences, top_scores.astype(float))] - } - return responder(results=results, error=200, message="Success") - - -def run_server(port=8008): - """This initializes the Tornad WSGI server to allow robust request handling. - - Parameters - ---------- - port : int, optional - Port number to serve the app on, by default 8008 - """ - - http_server = HTTPServer(WSGIContainer(app)) - http_server.listen(port) - - io_loop = IOLoop.instance() - tornado.autoreload.start(check_time=500) - print("Listening to port", port) - - try: - io_loop.start() - except KeyboardInterrupt: - pass - - -if __name__ == '__main__': - run_server(port=8008) diff --git a/text2vec/bin/utils.py b/text2vec/bin/utils.py deleted file mode 100644 index 57ac1ea..0000000 --- a/text2vec/bin/utils.py +++ /dev/null @@ -1,26 +0,0 @@ -import numpy as np - - -def log(message, **kwargs): - print(f"[INFO] {message}", flush=True, end=kwargs.get("end", "\n")) - - -def compute_angles(vectors): - """Computes the angles between vectors - - Parameters - ---------- - vectors : np.ndarray - (batch_size, embedding_size) - - Returns - ------- - np.ndarray - Cosine angles in degrees (batch_size, batch_size) - """ - - vectors /= np.linalg.norm(vectors, axis=1, keepdims=True) - cosine = np.dot(vectors, vectors.T) - cosine = np.clip(cosine, -1, 1) - degrees = np.arccos(cosine) * (180 / np.pi) - return degrees