File size: 5,582 Bytes
537011f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 | """HawkGPT 0.3 — Transformer Decoder with ALiBi positional encoding.
ALiBi: Attention with Linear Biases — no learned position embeddings.
Better for generalization, no position embedding limit.
"""
import math
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
import config
class MultiHeadSelfAttention(layers.Layer):
"""Multi-head self-attention with ALiBi positional biases."""
def __init__(self, embed_dim: int, num_heads: int, dropout: float = 0.05, **kwargs):
super().__init__(**kwargs)
assert embed_dim % num_heads == 0
self.num_heads = num_heads
self.head_dim = embed_dim // num_heads
self.qkv = layers.Dense(embed_dim * 3, use_bias=False)
self.out_proj = layers.Dense(embed_dim, use_bias=False)
self.dropout = layers.Dropout(dropout)
self.scale = math.sqrt(self.head_dim)
def call(self, x: tf.Tensor, training: bool = False) -> tf.Tensor:
B, T, C = tf.shape(x)[0], tf.shape(x)[1], tf.shape(x)[2]
qkv = self.qkv(x)
qkv = tf.reshape(qkv, (B, T, 3, self.num_heads, self.head_dim))
qkv = tf.transpose(qkv, (2, 0, 3, 1, 4))
q, k, v = qkv[0], qkv[1], qkv[2]
# Scaled dot-product attention
att = tf.matmul(q, tf.transpose(k, (0, 1, 3, 2))) / self.scale # (B, H, T, T)
# ALiBi: linear bias = -slope * |i - j|
# Slopes for each head: 2^(-8/num_heads * head_idx)
alibi = self._get_alibi(T, att.dtype)
att = att + alibi
# Causal mask
causal_mask = tf.linalg.band_part(tf.ones((T, T), dtype=att.dtype), -1, 0)
causal_mask = tf.reshape(causal_mask, (1, 1, T, T))
att = tf.where(tf.equal(causal_mask, 0), tf.constant(-1e9, dtype=att.dtype), att)
att = tf.nn.softmax(att, axis=-1)
att = self.dropout(att, training=training)
out = tf.matmul(att, v)
out = tf.transpose(out, (0, 2, 1, 3))
out = tf.reshape(out, (B, T, C))
return self.out_proj(out)
def _get_alibi(self, T: int, dtype) -> tf.Tensor:
"""Compute ALiBi bias matrix."""
slopes = []
for h in range(self.num_heads):
slopes.append(-2.0 ** (-8.0 * h / self.num_heads))
slopes = tf.constant(slopes, dtype=dtype) # (H,)
# Distance matrix: dist[i][j] = |i - j|
positions = tf.range(T, dtype=dtype)
dist = positions[:, None] - positions[None, :] # (T, T)
dist = tf.abs(dist)
# ALiBi: (H, T, T)
alibi = slopes[:, None, None] * dist[None, :, :]
return alibi
class FeedForward(layers.Layer):
def __init__(self, embed_dim: int, ff_dim: int, dropout: float = 0.05, **kwargs):
super().__init__(**kwargs)
self.net = keras.Sequential([
layers.Dense(ff_dim, activation="gelu"),
layers.Dense(embed_dim),
layers.Dropout(dropout),
])
def call(self, x: tf.Tensor, training: bool = False) -> tf.Tensor:
return self.net(x, training=training)
class TransformerBlock(layers.Layer):
def __init__(self, embed_dim, num_heads, ff_dim, dropout=0.05, **kwargs):
super().__init__(**kwargs)
self.ln1 = layers.LayerNormalization(epsilon=1e-5)
self.attn = MultiHeadSelfAttention(embed_dim, num_heads, dropout)
self.ln2 = layers.LayerNormalization(epsilon=1e-5)
self.ff = FeedForward(embed_dim, ff_dim, dropout)
self.drop = layers.Dropout(dropout)
def call(self, x, training=False):
x = x + self.drop(self.attn(self.ln1(x), training=training), training=training)
x = x + self.drop(self.ff(self.ln2(x), training=training), training=training)
return x
class GPTModel(keras.Model):
def __init__(
self,
vocab_size: int,
embed_dim: int = config.EMBED_DIM,
num_heads: int = config.NUM_HEADS,
num_layers: int = config.NUM_LAYERS,
ff_dim: int = config.FF_DIM,
max_seq_len: int = config.MAX_SEQ_LEN,
dropout: float = config.DROPOUT,
**kwargs,
):
super().__init__(**kwargs)
self.embed_dim = embed_dim
self.vocab_size = vocab_size
self.token_emb = layers.Embedding(vocab_size, embed_dim)
# NO positional embedding — ALiBi handles position via attention biases
self.drop = layers.Dropout(dropout)
self.blocks = [
TransformerBlock(embed_dim, num_heads, ff_dim, dropout)
for _ in range(num_layers)
]
self.ln_final = layers.LayerNormalization(epsilon=1e-5)
self.head = layers.Dense(vocab_size, use_bias=False)
def call(self, input_ids: tf.Tensor, training: bool = False) -> tf.Tensor:
B, T = tf.shape(input_ids)[0], tf.shape(input_ids)[1]
# Token embedding only — ALiBi adds position info in attention
x = self.token_emb(input_ids)
x = self.drop(x, training=training)
for block in self.blocks:
x = block(x, training=training)
x = self.ln_final(x)
logits = self.head(x)
return logits
def count_params(self) -> int:
return sum(tf.size(v).numpy() for v in self.trainable_variables)
def build_model(vocab_size: int) -> GPTModel:
model = GPTModel(vocab_size=vocab_size)
dummy = tf.zeros((1, config.MAX_SEQ_LEN), dtype=tf.int32)
model(dummy)
# Weight tying
model.head.kernel.assign(tf.transpose(model.token_emb.embeddings))
print(f"Model built: {model.count_params():,} parameters")
return model
|