Update app.py

app.py

@@ -21,6 +21,195 @@ print("🚀 Loading SAM-Z-1 Model...")
 MODEL_REPO = "Smilyai-labs/Sam-Z-1-tensorflow"
 CACHE_DIR = "./model_cache"
 
+# ============================================================================
+# Model Architecture Definitions (Required for Loading)
+# ============================================================================
+
+@tf.keras.saving.register_keras_serializable()
+class RotaryEmbedding(tf.keras.layers.Layer):
+    def __init__(self, dim, max_len=2048, theta=10000, **kwargs):
+        super().__init__(**kwargs)
+        self.dim = dim
+        self.max_len = max_len
+        self.theta = theta
+        self.built_cache = False
+    
+    def build(self, input_shape):
+        if not self.built_cache:
+            inv_freq = 1.0 / (self.theta ** (tf.range(0, self.dim, 2, dtype=tf.float32) / self.dim))
+            t = tf.range(self.max_len, dtype=tf.float32)
+            freqs = tf.einsum("i,j->ij", t, inv_freq)
+            emb = tf.concat([freqs, freqs], axis=-1)
+            
+            self.cos_cached = tf.constant(tf.cos(emb), dtype=tf.float32)
+            self.sin_cached = tf.constant(tf.sin(emb), dtype=tf.float32)
+            self.built_cache = True
+        
+        super().build(input_shape)
+    
+    def rotate_half(self, x):
+        x1, x2 = tf.split(x, 2, axis=-1)
+        return tf.concat([-x2, x1], axis=-1)
+    
+    def call(self, q, k):
+        seq_len = tf.shape(q)[2]
+        dtype = q.dtype
+        cos = tf.cast(self.cos_cached[:seq_len, :], dtype)[None, None, :, :]
+        sin = tf.cast(self.sin_cached[:seq_len, :], dtype)[None, None, :, :]
+        
+        q_rotated = (q * cos) + (self.rotate_half(q) * sin)
+        k_rotated = (k * cos) + (self.rotate_half(k) * sin)
+        
+        return q_rotated, k_rotated
+    
+    def get_config(self):
+        config = super().get_config()
+        config.update({"dim": self.dim, "max_len": self.max_len, "theta": self.theta})
+        return config
+
+
+@tf.keras.saving.register_keras_serializable()
+class RMSNorm(tf.keras.layers.Layer):
+    def __init__(self, epsilon=1e-5, **kwargs):
+        super().__init__(**kwargs)
+        self.epsilon = epsilon
+    
+    def build(self, input_shape):
+        self.scale = self.add_weight(name="scale", shape=(input_shape[-1],), initializer="ones")
+    
+    def call(self, x):
+        variance = tf.reduce_mean(tf.square(x), axis=-1, keepdims=True)
+        return x * tf.math.rsqrt(variance + self.epsilon) * self.scale
+    
+    def get_config(self):
+        config = super().get_config()
+        config.update({"epsilon": self.epsilon})
+        return config
+
+
+@tf.keras.saving.register_keras_serializable()
+class TransformerBlock(tf.keras.layers.Layer):
+    def __init__(self, d_model, n_heads, ff_dim, dropout, max_len, rope_theta, layer_idx=0, **kwargs):
+        super().__init__(**kwargs)
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.ff_dim = ff_dim
+        self.dropout_rate = dropout
+        self.max_len = max_len
+        self.rope_theta = rope_theta
+        self.head_dim = d_model // n_heads
+        self.layer_idx = layer_idx
+        
+        self.pre_attn_norm = RMSNorm()
+        self.pre_ffn_norm = RMSNorm()
+        
+        self.q_proj = tf.keras.layers.Dense(d_model, use_bias=False, name="q_proj")
+        self.k_proj = tf.keras.layers.Dense(d_model, use_bias=False, name="k_proj")
+        self.v_proj = tf.keras.layers.Dense(d_model, use_bias=False, name="v_proj")
+        self.out_proj = tf.keras.layers.Dense(d_model, use_bias=False, name="o_proj")
+        
+        self.rope = RotaryEmbedding(self.head_dim, max_len=max_len, theta=rope_theta)
+        
+        self.gate_proj = tf.keras.layers.Dense(ff_dim, use_bias=False, name="gate_proj")
+        self.up_proj = tf.keras.layers.Dense(ff_dim, use_bias=False, name="up_proj")
+        self.down_proj = tf.keras.layers.Dense(d_model, use_bias=False, name="down_proj")
+        
+        self.dropout = tf.keras.layers.Dropout(dropout)
+    
+    def call(self, x, training=None):
+        B, T, D = tf.shape(x)[0], tf.shape(x)[1], self.d_model
+        dtype = x.dtype
+        
+        # Attention
+        res = x
+        y = self.pre_attn_norm(x)
+        
+        q = tf.transpose(tf.reshape(self.q_proj(y), [B, T, self.n_heads, self.head_dim]), [0, 2, 1, 3])
+        k = tf.transpose(tf.reshape(self.k_proj(y), [B, T, self.n_heads, self.head_dim]), [0, 2, 1, 3])
+        v = tf.transpose(tf.reshape(self.v_proj(y), [B, T, self.n_heads, self.head_dim]), [0, 2, 1, 3])
+        
+        q, k = self.rope(q, k)
+        
+        scores = tf.matmul(q, k, transpose_b=True) / tf.sqrt(tf.cast(self.head_dim, dtype))
+        
+        mask = tf.where(
+            tf.linalg.band_part(tf.ones([T, T], dtype=dtype), -1, 0) == 0,
+            tf.constant(-1e9, dtype=dtype),
+            tf.constant(0.0, dtype=dtype)
+        )
+        scores += mask
+        attn = tf.matmul(tf.nn.softmax(scores, axis=-1), v)
+        
+        attn = tf.reshape(tf.transpose(attn, [0, 2, 1, 3]), [B, T, D])
+        x = res + self.dropout(self.out_proj(attn), training=training)
+        
+        # FFN (SwiGLU)
+        res = x
+        y = self.pre_ffn_norm(x)
+        ffn = self.down_proj(tf.keras.activations.silu(self.gate_proj(y)) * self.up_proj(y))
+        
+        return res + self.dropout(ffn, training=training)
+    
+    def get_config(self):
+        config = super().get_config()
+        config.update({
+            "d_model": self.d_model,
+            "n_heads": self.n_heads,
+            "ff_dim": self.ff_dim,
+            "dropout": self.dropout_rate,
+            "max_len": self.max_len,
+            "rope_theta": self.rope_theta,
+            "layer_idx": self.layer_idx
+        })
+        return config
+
+
+@tf.keras.saving.register_keras_serializable()
+class SAM1Model(tf.keras.Model):
+    def __init__(self, **kwargs):
+        super().__init__()
+        if 'config' in kwargs and isinstance(kwargs['config'], dict):
+            self.cfg = kwargs['config']
+        elif 'vocab_size' in kwargs:
+            self.cfg = kwargs
+        else:
+            self.cfg = kwargs.get('cfg', kwargs)
+        
+        self.embed = tf.keras.layers.Embedding(self.cfg['vocab_size'], self.cfg['d_model'], name="embed_tokens")
+        
+        ff_dim = int(self.cfg['d_model'] * self.cfg['ff_mult'])
+        block_args = {
+            'd_model': self.cfg['d_model'],
+            'n_heads': self.cfg['n_heads'],
+            'ff_dim': ff_dim,
+            'dropout': self.cfg['dropout'],
+            'max_len': self.cfg['max_len'],
+            'rope_theta': self.cfg['rope_theta']
+        }
+        
+        self.blocks = []
+        for i in range(self.cfg['n_layers']):
+            block = TransformerBlock(name=f"block_{i}", layer_idx=i, **block_args)
+            self.blocks.append(block)
+        
+        self.norm = RMSNorm(name="final_norm")
+        self.lm_head = tf.keras.layers.Dense(self.cfg['vocab_size'], use_bias=False, name="lm_head")
+    
+    def call(self, input_ids, training=None):
+        x = self.embed(input_ids)
+        
+        for block in self.blocks:
+            x = block(x, training=training)
+        
+        return self.lm_head(self.norm(x))
+    
+    def get_config(self):
+        base_config = super().get_config()
+        base_config['config'] = self.cfg
+        return base_config
+
+print("✅ Model architecture registered")
+
 # Download model files
 config_path = hf_hub_download(MODEL_REPO, "config.json", cache_dir=CACHE_DIR)
 model_path = hf_hub_download(MODEL_REPO, "model.keras", cache_dir=CACHE_DIR)