File size: 2,518 Bytes
47dfda4 ac2b9c4 47dfda4 ac2b9c4 47dfda4 | 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 | import tensorflow as tf
from keras import layers
class PositionEmbedding(layers.Layer):
def __init__(self, block_size, **kwargs):
super().__init__(**kwargs)
self.block_size = block_size
def build(self, input_shape):
self.pos_emb = self.add_weight(
name="pos_emb", shape=(self.block_size, input_shape[-1]), initializer="random_normal"
)
def call(self, x):
T = tf.shape(x)[1]
return x + self.pos_emb[tf.newaxis, :T, :]
def get_config(self):
config = super().get_config()
config["block_size"] = self.block_size
return config
def build_gpt(vocab_size, block_size, n_layer=4, n_head=4, n_embd=128, dropout=0.1, name="indigo"):
tokens = tf.keras.Input(shape=(None,), dtype="int64", name="tokens")
x = layers.Embedding(vocab_size, n_embd, name="tok_emb")(tokens)
x = PositionEmbedding(block_size, name="pos_emb")(x)
x = layers.Dropout(dropout)(x)
for i in range(n_layer):
xn = layers.LayerNormalization(epsilon=1e-5, name=f"ln1_{i}")(x)
attn = layers.MultiHeadAttention(
num_heads=n_head, key_dim=n_embd // n_head, dropout=dropout, name=f"attn_{i}"
)
x = x + attn(xn, xn, use_causal_mask=True)
xn = layers.LayerNormalization(epsilon=1e-5, name=f"ln2_{i}")(x)
h = layers.Dense(4 * n_embd, activation="gelu", name=f"fc_{i}")(xn)
h = layers.Dense(n_embd, name=f"proj_{i}")(h)
h = layers.Dropout(dropout)(h)
x = x + h
x = layers.LayerNormalization(epsilon=1e-5, name="ln_f")(x)
logits = layers.Dense(vocab_size, use_bias=False, name="head")(x)
return tf.keras.Model(tokens, logits, name=name)
@tf.function(reduce_retracing=True)
def _langkah(model, idx_cond, temperature, top_k):
logits = model(idx_cond, training=False)[:, -1, :]
logits = logits / max(temperature, 1e-8)
if top_k is not None:
k = min(top_k, int(logits.shape[-1]))
vals, _ = tf.math.top_k(logits, k=k)
logits = tf.where(
logits < vals[:, -1:],
tf.fill(tf.shape(logits), tf.float32.min),
logits,
)
return tf.random.categorical(logits, num_samples=1, dtype=tf.int64)
def generate(model, idx, max_new_tokens, block_size, temperature=1.0, top_k=None):
for _ in range(max_new_tokens):
idx_cond = idx[:, -block_size:]
next_id = _langkah(model, idx_cond, temperature, top_k)
idx = tf.concat([idx, next_id], axis=1)
return idx
|