model.py

import tensorflow as tf
import numpy as np

class SquaredReLU(tf.keras.layers.Layer):
    def __init__(self, **kwargs):
        super(SquaredReLU, self).__init__(**kwargs)

    def call(self, x):
        return tf.square(tf.nn.relu(x))

class TokenShift(tf.keras.layers.Layer):
    """
    Position encoding via Token Shifting.
    Shifts the sequence to provide local positional context without additive embeddings.
    """
    def __init__(self, shift_amount=1, **kwargs):
        super(TokenShift, self).__init__(**kwargs)
        self.shift_amount = shift_amount

    def call(self, x):
        # x shape: [batch, seq_len, dim]
        seq_len = tf.shape(x)[1]
        shifted = tf.roll(x, shift=self.shift_amount, axis=1)
        
        # Create a mask to zero out the elements that wrapped around
        # The first 'shift_amount' elements are the ones that wrapped from the end
        mask = tf.ones([seq_len], dtype=x.dtype)
        mask = tf.concat([tf.zeros([self.shift_amount], dtype=x.dtype), mask[self.shift_amount:]], axis=0)
        mask = tf.reshape(mask, [1, seq_len, 1])
        
        return x + (shifted * mask)

class GroupNorm(tf.keras.layers.Layer):
    """
    Group Normalization instead of LayerNorm.
    """
    def __init__(self, groups=32, eps=1e-6, **kwargs):
        super(GroupNorm, self).__init__(**kwargs)
        self.groups = groups
        self.eps = eps

    def build(self, input_shape):
        dim = input_shape[-1]
        if dim % self.groups != 0:
            raise ValueError(f"Dimension {dim} must be divisible by groups {self.groups}")
        self.gamma = self.add_weight(shape=(dim,), initializer='ones', trainable=True, name='gamma')
        self.beta = self.add_weight(shape=(dim,), initializer='zeros', trainable=True, name='beta')

    def call(self, x):
        # x shape: [batch, seq_len, dim]
        batch, seq_len, dim = tf.shape(x)[0], tf.shape(x)[1], tf.shape(x)[2]
        x = tf.reshape(x, [batch, seq_len, self.groups, dim // self.groups])
        
        mean = tf.reduce_mean(x, axis=-1, keepdims=True)
        var = tf.reduce_mean(tf.square(x - mean), axis=-1, keepdims=True)
        
        x = (x - mean) / tf.sqrt(var + self.eps)
        x = tf.reshape(x, [batch, seq_len, dim])
        
        return x * self.gamma + self.beta

class StochasticDepth(tf.keras.layers.Layer):
    """
    Regularization: Randomly drops layers during training.
    """
    def __init__(self, drop_prob=0.1, **kwargs):
        super(StochasticDepth, self).__init__(**kwargs)
        self.drop_prob = drop_prob

    def call(self, x, training=False):
        if not training or self.drop_prob == 0:
            return x
        
        keep_prob = 1.0 - self.drop_prob
        random_tensor = tf.random.uniform(shape=[tf.shape(x)[0]], minval=0, maxval=1)
        binary_mask = tf.cast(random_tensor < keep_prob, tf.float32)
        
        # Reshape mask for broadcasting: [batch, 1, 1]
        mask = tf.reshape(binary_mask, [-1, 1, 1])
        
        return x * mask / keep_prob

class TimeMix(tf.keras.layers.Layer):
    """
    Unique sequence mixing mechanism for Tera v3.
    Uses a depth-wise temporal convolution to mix information across time.
    """
    def __init__(self, dim, **kwargs):
        super(TimeMix, self).__init__(**kwargs)
        self.dim = dim
        # Depth-wise convolution for temporal mixing
        self.conv = tf.keras.layers.DepthwiseConv1D(
            kernel_size=3, padding='same', activation=None
        )
        self.norm = GroupNorm()
        self.proj = tf.keras.layers.Dense(dim)

    def call(self, x):
        # x shape: [batch, seq_len, dim]
        residual = x
        x = self.norm(x)
        x = self.conv(x)
        x = self.proj(x)
        return x + residual

class ChannelMix(tf.keras.layers.Layer):
    """
    Feed-forward network using Squared ReLU.
    """
    def __init__(self, dim, expand_factor=4, **kwargs):
        super(ChannelMix, self).__init__(**kwargs)
        self.dim = dim
        self.inner_dim = dim * expand_factor
        self.norm = GroupNorm()
        self.fc1 = tf.keras.layers.Dense(self.inner_dim)
        self.sq_relu = SquaredReLU()
        self.fc2 = tf.keras.layers.Dense(dim)

    def call(self, x):
        residual = x
        x = self.norm(x)
        x = self.fc1(x)
        x = self.sq_relu(x)
        x = self.fc2(x)
        return x + residual

class TeraBlock(tf.keras.layers.Layer):
    """
    The core block of Tera v3.
    """
    def __init__(self, dim, drop_prob=0.1, **kwargs):
        super(TeraBlock, self).__init__(**kwargs)
        self.token_shift = TokenShift()
        self.time_mix = TimeMix(dim)
        self.channel_mix = ChannelMix(dim)
        self.stoch_depth = StochasticDepth(drop_prob)

    def call(self, x, training=False):
        # Positional encoding via shift
        x = self.token_shift(x)
        
        # Sequence mixing
        x = x + self.stoch_depth(self.time_mix(x), training=training)
        
        # Feed-forward mixing
        x = x + self.stoch_depth(self.channel_mix(x), training=training)
        
        return x

class TeraVisionEncoder(tf.keras.layers.Layer):
    """
    Unique Vision Capability: Spectral-Spatial Shift Projection.
    Instead of standard ViT, it uses a custom patch-shifting mechanism.
    """
    def __init__(self, patch_size=16, embed_dim=512, **kwargs):
        super(TeraVisionEncoder, self).__init__(**kwargs)
        self.patch_size = patch_size
        self.embed_dim = embed_dim
        
        # Use a custom convolution for patch embedding to avoid 'Google' style ViT layers
        self.patch_embed = tf.keras.layers.Conv2D(
            filters=embed_dim, 
            kernel_size=patch_size, 
            strides=patch_size, 
            padding='valid'
        )
        
    def call(self, images):
        # images: [batch, h, w, c]
        x = self.patch_embed(images) # [batch, h/p, w/p, dim]
        # Flatten spatial dimensions to sequence
        batch = tf.shape(x)[0]
        x = tf.reshape(x, [batch, -1, self.embed_dim])
        return x

class TeraV3(tf.keras.Model):
    """
    Tera v3 Language Model.
    """
    def __init__(self, vocab_size, dim=512, depth=12, embed_dim=512, **kwargs):
        super(TeraV3, self).__init__(**kwargs)
        
        # Untied Embeddings: Input and Output embeddings are separate
        self.input_embedding = tf.keras.layers.Embedding(vocab_size, embed_dim)
        self.output_projection = tf.keras.layers.Dense(vocab_size)
        
        # Vision Encoder
        self.vision_encoder = TeraVisionEncoder(embed_dim=dim)
        
        # Core Architecture
        self.blocks = [TeraBlock(dim) for _ in range(depth)]
        self.final_norm = GroupNorm()
        
        # Project input embedding to model dim if they differ
        if embed_dim != dim:
            self.input_proj = tf.keras.layers.Dense(dim)
        else:
            self.input_proj = tf.keras.layers.Lambda(lambda x: x)

    def call(self, inputs, training=False, vision_inputs=None):
        # Handle vision inputs
        if vision_inputs is not None:
            v_emb = self.vision_encoder(vision_inputs)
            # Concat vision tokens with text tokens
            x = self.input_embedding(inputs)
            x = self.input_proj(x)
            x = tf.concat([v_emb, x], axis=1)
        else:
            x = self.input_embedding(inputs)
            x = self.input_proj(x)
        
        # Pass through Tera Blocks
        for block in self.blocks:
            x = block(x, training=training)
            
        x = self.final_norm(x)
        
        # Project to vocab for text generation
        # For vision-text models, usually only the text part produces logits
        # but for simplicity, we project the whole sequence.
        logits = self.output_projection(x)
        
        return logits

# Example instantiation
if __name__ == "__main__":
    vocab_size = 10000
    model = TeraV3(vocab_size=vocab_size)
    
    # Dummy text input
    text_in = tf.constant([[1, 2, 3, 4, 5]], dtype=tf.int32)
    # Dummy vision input
    vision_in = tf.random.uniform([1, 224, 224, 3])
    
    output = model(text_in, vision_inputs=vision_in)
    print("Output shape:", output.shape) # [batch, (num_patches + seq_len), vocab_size]