Spaces:
Sleeping
Sleeping
File size: 3,766 Bytes
9573dc0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 | """
Custom Keras layers for BERT metagenome model.
These layers must be registered as custom_objects when loading the model.
"""
import tensorflow as tf
@tf.keras.utils.register_keras_serializable(package="deepG")
class layer_pos_embedding(tf.keras.layers.Layer):
"""Token + Positional Embedding layer for BERT."""
def __init__(self, maxlen=1000, vocabulary_size=6, embed_dim=600, **kwargs):
super().__init__(**kwargs)
self.maxlen = int(maxlen)
self.vocabulary_size = int(vocabulary_size)
self.embed_dim = int(embed_dim)
self.token_emb = tf.keras.layers.Embedding(
input_dim=self.vocabulary_size,
output_dim=self.embed_dim,
name="token_emb",
)
self.pos_emb = tf.keras.layers.Embedding(
input_dim=self.maxlen,
output_dim=self.embed_dim,
name="pos_emb",
)
def call(self, x):
x = tf.cast(x, tf.int32)
L = tf.shape(x)[1]
positions = tf.range(start=0, limit=L, delta=1)
positions = self.pos_emb(positions)
tokens = self.token_emb(x)
return tokens + positions
def get_config(self):
cfg = super().get_config()
cfg.update(
dict(
maxlen=self.maxlen,
vocabulary_size=self.vocabulary_size,
embed_dim=self.embed_dim,
)
)
return cfg
@tf.keras.utils.register_keras_serializable(package="deepG")
class layer_transformer_block(tf.keras.layers.Layer):
"""Transformer block with Multi-Head Attention and Feed-Forward Network."""
def __init__(
self,
num_heads=16,
head_size=250,
dropout_rate=0.0,
ff_dim=2400.0,
vocabulary_size=6,
embed_dim=600,
**kwargs
):
super().__init__(**kwargs)
self.num_heads = int(num_heads)
self.head_size = int(head_size)
self.dropout_rate = float(dropout_rate)
self.ff_dim = int(ff_dim)
self.vocabulary_size = int(vocabulary_size)
self.embed_dim = int(embed_dim)
self.mha = tf.keras.layers.MultiHeadAttention(
num_heads=self.num_heads,
key_dim=self.head_size,
dropout=self.dropout_rate,
name="mha",
)
self.ffn1 = tf.keras.layers.Dense(self.ff_dim, activation=tf.nn.gelu, name="ffn1")
self.ffn2 = tf.keras.layers.Dense(self.embed_dim, name="ffn2")
self.ln1 = tf.keras.layers.LayerNormalization(epsilon=1e-6, name="ln1")
self.ln2 = tf.keras.layers.LayerNormalization(epsilon=1e-6, name="ln2")
self.drop1 = tf.keras.layers.Dropout(self.dropout_rate, name="drop1")
self.drop2 = tf.keras.layers.Dropout(self.dropout_rate, name="drop2")
def call(self, x, training=False):
attn = self.mha(x, x, training=training)
attn = self.drop1(attn, training=training)
x = x + attn
x = self.ln1(x)
f = self.ffn2(self.ffn1(x))
f = self.drop2(f, training=training)
x = x + f
x = self.ln2(x)
return x
def get_config(self):
cfg = super().get_config()
cfg.update(
dict(
num_heads=self.num_heads,
head_size=self.head_size,
dropout_rate=self.dropout_rate,
ff_dim=self.ff_dim,
vocabulary_size=self.vocabulary_size,
embed_dim=self.embed_dim,
)
)
return cfg
def get_custom_objects():
"""Return dictionary of custom objects needed for model loading."""
return {
"layer_pos_embedding": layer_pos_embedding,
"layer_transformer_block": layer_transformer_block,
}
|