Update model.py

model.py

@@ -2,106 +2,13 @@ import tensorflow as tf
 from tensorflow import keras
 from tensorflow.keras import layers
 import numpy as np
-from typing import Optional
-
-class PositionalEncoding(layers.Layer):
-    """Positional encoding layer for transformer"""
-    
-    def __init__(self, max_length: int, d_model: int, **kwargs):
-        super().__init__(**kwargs)
-        self.max_length = max_length
-        self.d_model = d_model
-        
-        # Create positional encoding matrix
-        position = np.arange(max_length)[:, np.newaxis]
-        div_term = np.exp(np.arange(0, d_model, 2) * -(np.log(10000.0) / d_model))
-        
-        pe = np.zeros((max_length, d_model))
-        pe[:, 0::2] = np.sin(position * div_term)
-        pe[:, 1::2] = np.cos(position * div_term)
-        
-        self.positional_encoding = tf.constant(pe, dtype=tf.float32)
-    
-    def call(self, x):
-        seq_length = tf.shape(x)[1]
-        return x + self.positional_encoding[:seq_length, :]
-    
-    def get_config(self):
-        config = super().get_config()
-        config.update({
-            'max_length': self.max_length,
-            'd_model': self.d_model
-        })
-        return config
-
-
-class TransformerBlock(layers.Layer):
-    """Transformer decoder block"""
-    
-    def __init__(self, d_model: int, num_heads: int, ff_dim: int, 
-                 dropout_rate: float = 0.1, **kwargs):
-        super().__init__(**kwargs)
-        self.d_model = d_model
-        self.num_heads = num_heads
-        self.ff_dim = ff_dim
-        self.dropout_rate = dropout_rate
-        
-        self.attention = layers.MultiHeadAttention(
-            num_heads=num_heads, 
-            key_dim=d_model // num_heads,
-            dropout=dropout_rate
-        )
-        self.ffn = keras.Sequential([
-            layers.Dense(ff_dim, activation='gelu'),
-            layers.Dropout(dropout_rate),
-            layers.Dense(d_model),
-            layers.Dropout(dropout_rate)
-        ])
-        self.layernorm1 = layers.LayerNormalization(epsilon=1e-6)
-        self.layernorm2 = layers.LayerNormalization(epsilon=1e-6)
-        self.dropout = layers.Dropout(dropout_rate)
-    
-    def call(self, x, training=False, mask=None):
-        # Causal self-attention
-        attn_output = self.attention(
-            query=x, 
-            value=x, 
-            key=x,
-            attention_mask=mask,
-            training=training
-        )
-        attn_output = self.dropout(attn_output, training=training)
-        out1 = self.layernorm1(x + attn_output)
-        
-        # Feed forward network
-        ffn_output = self.ffn(out1, training=training)
-        return self.layernorm2(out1 + ffn_output)
-    
-    def get_config(self):
-        config = super().get_config()
-        config.update({
-            'd_model': self.d_model,
-            'num_heads': self.num_heads,
-            'ff_dim': self.ff_dim,
-            'dropout_rate': self.dropout_rate
-        })
-        return config
-
 
 class VedaProgrammingLLM(keras.Model):
     """Veda Programming Language Model"""
     
-    def __init__(
-        self,
-        vocab_size: int,
-        max_length: int = 512,
-        d_model: int = 256,
-        num_heads: int = 8,
-        num_layers: int = 6,
-        ff_dim: int = 1024,
-        dropout_rate: float = 0.1,
-        **kwargs
-    ):
+    def __init__(self, vocab_size: int, max_length: int = 256, 
+                 d_model: int = 128, num_heads: int = 4, 
+                 num_layers: int = 2, ff_dim: int = 256, **kwargs):
         super().__init__(**kwargs)
         
         self.vocab_size = vocab_size
@@ -110,117 +17,85 @@ class VedaProgrammingLLM(keras.Model):
         self.num_heads = num_heads
         self.num_layers = num_layers
         self.ff_dim = ff_dim
-        self.dropout_rate = dropout_rate
         
-        # Embedding layers
-        self.token_embedding = layers.Embedding(
-            input_dim=vocab_size, 
-            output_dim=d_model
-        )
-        self.positional_encoding = PositionalEncoding(max_length, d_model)
-        self.dropout = layers.Dropout(dropout_rate)
+        # Embeddings
+        self.token_embedding = layers.Embedding(vocab_size, d_model)
+        self.pos_embedding = layers.Embedding(max_length, d_model)
+        self.dropout = layers.Dropout(0.1)
         
-        # Transformer blocks
-        self.transformer_blocks = [
-            TransformerBlock(d_model, num_heads, ff_dim, dropout_rate)
-            for _ in range(num_layers)
-        ]
+        # Transformer layers
+        self.transformer_blocks = []
+        for _ in range(num_layers):
+            self.transformer_blocks.append({
+                'attn': layers.MultiHeadAttention(num_heads=num_heads, key_dim=d_model // num_heads),
+                'ffn': keras.Sequential([
+                    layers.Dense(ff_dim, activation='relu'),
+                    layers.Dense(d_model)
+                ]),
+                'ln1': layers.LayerNormalization(),
+                'ln2': layers.LayerNormalization(),
+                'dropout1': layers.Dropout(0.1),
+                'dropout2': layers.Dropout(0.1)
+            })
         
-        # Output layer
+        self.final_ln = layers.LayerNormalization()
         self.output_layer = layers.Dense(vocab_size)
     
-    def _create_causal_mask(self, seq_length):
-        """Create causal attention mask"""
-        mask = tf.linalg.band_part(
-            tf.ones((seq_length, seq_length)), -1, 0
-        )
-        return mask
-    
     def call(self, inputs, training=False):
-        seq_length = tf.shape(inputs)[1]
+        seq_len = tf.shape(inputs)[1]
         
         # Create causal mask
-        mask = self._create_causal_mask(seq_length)
+        mask = self._create_causal_mask(seq_len)
         
         # Embeddings
-        x = self.token_embedding(inputs)
-        x = x * tf.math.sqrt(tf.cast(self.d_model, tf.float32))
-        x = self.positional_encoding(x)
+        positions = tf.range(seq_len)
+        x = self.token_embedding(inputs) + self.pos_embedding(positions)
         x = self.dropout(x, training=training)
         
         # Transformer blocks
-        for transformer_block in self.transformer_blocks:
-            x = transformer_block(x, training=training, mask=mask)
+        for block in self.transformer_blocks:
+            # Self attention
+            attn_out = block['attn'](x, x, attention_mask=mask, training=training)
+            attn_out = block['dropout1'](attn_out, training=training)
+            x = block['ln1'](x + attn_out)
+            
+            # FFN
+            ffn_out = block['ffn'](x)
+            ffn_out = block['dropout2'](ffn_out, training=training)
+            x = block['ln2'](x + ffn_out)
         
-        # Output projection
-        logits = self.output_layer(x)
-        return logits
+        x = self.final_ln(x)
+        return self.output_layer(x)
+    
+    def _create_causal_mask(self, seq_len):
+        """Create causal attention mask"""
+        mask = tf.linalg.band_part(tf.ones((seq_len, seq_len)), -1, 0)
+        return mask
     
-    def generate(
-        self, 
-        prompt_tokens: list, 
-        max_new_tokens: int = 100,
-        temperature: float = 0.7,
-        top_k: int = 50,
-        top_p: float = 0.9
-    ):
-        """Generate code given a prompt"""
+    def generate(self, prompt_tokens: list, max_new_tokens: int = 50,
+                 temperature: float = 0.8, top_k: int = 40):
+        """Generate code"""
         generated = list(prompt_tokens)
         
         for _ in range(max_new_tokens):
-            # Truncate if too long
             context = generated[-self.max_length:]
+            input_tensor = tf.constant([context], dtype=tf.int32)
             
-            # Get predictions
-            input_tensor = tf.expand_dims(context, 0)
             logits = self(input_tensor, training=False)
-            next_token_logits = logits[0, -1, :] / temperature
+            next_logits = logits[0, -1, :] / temperature
             
-            # Apply top-k filtering
+            # Top-k sampling
             if top_k > 0:
-                top_k_logits, top_k_indices = tf.math.top_k(
-                    next_token_logits, k=min(top_k, self.vocab_size)
-                )
-                # Create mask for non-top-k tokens
-                indices_to_remove = tf.less(
-                    next_token_logits, 
-                    top_k_logits[-1]
-                )
-                next_token_logits = tf.where(
-                    indices_to_remove,
-                    tf.ones_like(next_token_logits) * float('-inf'),
-                    next_token_logits
-                )
-            
-            # Apply top-p (nucleus) filtering
-            if top_p < 1.0:
-                sorted_logits = tf.sort(next_token_logits, direction='DESCENDING')
-                sorted_probs = tf.nn.softmax(sorted_logits)
-                cumulative_probs = tf.cumsum(sorted_probs)
-                
-                # Find cutoff
-                sorted_indices_to_remove = cumulative_probs > top_p
-                sorted_indices_to_remove = tf.concat([
-                    [False], 
-                    sorted_indices_to_remove[:-1]
-                ], axis=0)
-                
-                sorted_logits = tf.where(
-                    sorted_indices_to_remove,
-                    tf.ones_like(sorted_logits) * float('-inf'),
-                    sorted_logits
-                )
+                top_k_logits, top_k_indices = tf.math.top_k(next_logits, k=min(top_k, self.vocab_size))
+                probs = tf.nn.softmax(top_k_logits)
+                idx = tf.random.categorical(tf.expand_dims(tf.math.log(probs + 1e-10), 0), 1)[0, 0]
+                next_token = top_k_indices[idx].numpy()
+            else:
+                probs = tf.nn.softmax(next_logits)
+                next_token = tf.random.categorical(tf.expand_dims(tf.math.log(probs + 1e-10), 0), 1)[0, 0].numpy()
             
-            # Sample from distribution
-            probs = tf.nn.softmax(next_token_logits)
-            next_token = tf.random.categorical(
-                tf.expand_dims(next_token_logits, 0), 
-                num_samples=1
-            )[0, 0]
+            generated.append(int(next_token))
             
-            generated.append(int(next_token.numpy()))
-            
-            # Stop if end token
             if next_token == 3:  # END token
                 break
         
@@ -233,49 +108,5 @@ class VedaProgrammingLLM(keras.Model):
             'd_model': self.d_model,
             'num_heads': self.num_heads,
             'num_layers': self.num_layers,
-            'ff_dim': self.ff_dim,
-            'dropout_rate': self.dropout_rate
-        }
-    
-    @classmethod
-    def from_config(cls, config):
-        return cls(**config)
-
-
-def create_veda_model(
-    vocab_size: int,
-    max_length: int = 512,
-    model_size: str = "small"
-) -> VedaProgrammingLLM:
-    """Factory function to create Veda Programming model"""
-    
-    configs = {
-        "small": {
-            "d_model": 256,
-            "num_heads": 4,
-            "num_layers": 4,
-            "ff_dim": 512
-        },
-        "medium": {
-            "d_model": 512,
-            "num_heads": 8,
-            "num_layers": 6,
-            "ff_dim": 1024
-        },
-        "large": {
-            "d_model": 768,
-            "num_heads": 12,
-            "num_layers": 12,
-            "ff_dim": 2048
-        }
-    }
-    
-    config = configs.get(model_size, configs["small"])
-    
-    model = VedaProgrammingLLM(
-        vocab_size=vocab_size,
-        max_length=max_length,
-        **config
-    )
-    
-    return model
+            'ff_dim': self.ff_dim
+        }