Update api.py

api.py

@@ -84,61 +84,67 @@ class SwiGLU(tf.keras.layers.Layer):
         x_val, x_gate = tf.split(x_proj, 2, axis=-1)    
         return self.out(x_val * tf.nn.silu(x_gate))    
             
-class GPTBlock(tf.keras.layers.Layer):    
-    def __init__(self, d_model, d_ff, num_heads=8, dropout_rate=0.1, adapter_dim=64):      
-        super().__init__()      
-        self.ln1 = tf.keras.layers.LayerNormalization(epsilon=1e-5)      
-        self.mha = tf.keras.layers.MultiHeadAttention(num_heads=num_heads, key_dim=d_model // num_heads)      
-        self.dropout1 = tf.keras.layers.Dropout(dropout_rate)     
-        self.adapter_down = tf.keras.layers.Dense(adapter_dim, activation='gelu')     
-        self.adapter_up = tf.keras.layers.Dense(d_model)      
-      
-        self.ln2 = tf.keras.layers.LayerNormalization(epsilon=1e-5)      
-        self.ffn = SwiGLU(d_model, d_ff)      
-        self.dropout2 = tf.keras.layers.Dropout(dropout_rate)     
-        self.rope = RotaryPositionalEmbedding(d_model // num_heads)      
-      
-    def call(self, x, training=False):      
-        x_norm = self.ln1(x)      
-        b, s, _ = tf.shape(x_norm)[0], tf.shape(x_norm)[1], tf.shape(x_norm)[2]      
-        h = self.mha.num_heads      
-        d = x_norm.shape[-1] // h      
-      
-        qkv = tf.reshape(x_norm, [b, s, h, d])      
-        qkv = tf.transpose(qkv, [0, 2, 1, 3])      
-        q = self.rope(qkv)      
-        k = self.rope(qkv)      
-        q = tf.reshape(tf.transpose(q, [0, 2, 1, 3]), [b, s, h * d])      
-        k = tf.reshape(tf.transpose(k, [0, 2, 1, 3]), [b, s, h * d])      
-      
-        attn_out = self.mha(query=q, value=x_norm, key=k, use_causal_mask=True, training=training)      
-        attn_out = self.dropout1(attn_out, training=training)      
+class Block(tf.keras.layers.Layer):  
+    def __init__(self, d_model, d_ff, num_heads=8, dropout_rate=0.05, adapter_dim=64):    
+        super().__init__()    
+        self.ln1 = tf.keras.layers.LayerNormalization(epsilon=1e-5)    
+        self.mha = tf.keras.layers.MultiHeadAttention(num_heads=num_heads, key_dim=d_model // num_heads)    
+        self.dropout1 = tf.keras.layers.Dropout(dropout_rate)   
+        self.adapter_down = tf.keras.layers.Dense(adapter_dim, activation='gelu')   
+        self.adapter_up = tf.keras.layers.Dense(d_model)    
     
-        adapter_out = self.adapter_up(self.adapter_down(attn_out))    
-        attn_out = attn_out + adapter_out      
-      
-        x = x + attn_out      
-        ffn_out = self.ffn(self.ln2(x))      
-        x = x + self.dropout2(ffn_out, training=training)      
-        return x    
+        self.ln2 = tf.keras.layers.LayerNormalization(epsilon=1e-5)    
+        self.ffn = SwiGLU(d_model, d_ff)    
+        self.dropout2 = tf.keras.layers.Dropout(dropout_rate)   
+        self.rope = RotaryPositionalEmbedding(d_model // num_heads)    
     
-class InteractGPT(tf.keras.Model):      
-    def __init__(self, vocab_size, seq_len, d_model, d_ff, n_layers, num_heads=8, dropout_rate=0.1):      
-        super().__init__()      
-        self.token_embedding = tf.keras.layers.Embedding(vocab_size, d_model)      
-        self.blocks = [GPTBlock(d_model, d_ff, num_heads, dropout_rate) for _ in range(n_layers)]      
-        self.ln_f = tf.keras.layers.LayerNormalization(epsilon=1e-5)      
-      
-    def call(self, x, training=False):      
-        x = self.token_embedding(x)      
-        for block in self.blocks:      
-            x = block(x, training=training)      
-        x = self.ln_f(x)      
-        logits = tf.matmul(x, self.token_embedding.embeddings, transpose_b=True)      
-        return logits      
+    def call(self, x, training=False):    
+        x_norm = self.ln1(x)    
+        b, s, _ = tf.shape(x_norm)[0], tf.shape(x_norm)[1], tf.shape(x_norm)[2]    
+        h = self.mha.num_heads    
+        d = x_norm.shape[-1] // h    
     
-model = InteractGPT(vocab_size=vocab_size, seq_len=max_len, d_model=256, d_ff=1024, n_layers=6)        
-      
+        qkv = tf.reshape(x_norm, [b, s, h, d])    
+        qkv = tf.transpose(qkv, [0, 2, 1, 3])    
+        q = self.rope(qkv)    
+        k = self.rope(qkv)    
+        q = tf.reshape(tf.transpose(q, [0, 2, 1, 3]), [b, s, h * d])    
+        k = tf.reshape(tf.transpose(k, [0, 2, 1, 3]), [b, s, h * d])    
+    
+        attn_out = self.mha(query=q, value=x_norm, key=k, use_causal_mask=True, training=training)    
+        attn_out = self.dropout1(attn_out, training=training)    
+  
+        adapter_out = self.adapter_up(self.adapter_down(attn_out))  
+        attn_out = attn_out + adapter_out    
+    
+        x = x + attn_out    
+        ffn_out = self.ffn(self.ln2(x))    
+        x = x + self.dropout2(ffn_out, training=training)    
+        return x
+    
+class Flexi(tf.keras.Model):  
+    def __init__(self, vocab_size, seq_len, d_model, d_ff, n_layers, num_heads=8, dropout_rate=0.05):  
+        super().__init__()  
+        self.token_embedding = tf.keras.layers.Embedding(vocab_size, d_model)  
+        self.blocks = [Block(d_model, d_ff, num_heads, dropout_rate) for _ in range(n_layers)]  
+        self.ln_f = tf.keras.layers.LayerNormalization(epsilon=1e-5)  
+  
+    def call(self, x, training=False):  
+        x = self.token_embedding(x)  
+        for block in self.blocks:  
+            x = block(x, training=training)  
+        x = self.ln_f(x)  
+        logits = tf.matmul(x, self.token_embedding.embeddings, transpose_b=True)  
+        return logits     
+    
+model = Flexi(
+    vocab_size=vocab_size,
+    seq_len=max_len,
+    d_model=256,
+    d_ff=1024,       
+    n_layers=16
+)
+
 dummy_input = tf.zeros((1, max_len), dtype=tf.int32)  # 배치1, 시퀀스길이 max_len      
 _ = model(dummy_input)  # 모델이 빌드됨      
 model.load_weights("Flexi.weights.h5")