codon-model 0.0.1__tar.gz → 0.0.2__tar.gz

{codon_model-0.0.1/codon_model.egg-info → codon_model-0.0.2}/PKG-INFO

@@ -1,15 +1,18 @@
 Metadata-Version: 2.4
 Name: codon-model
-Version: 0.0.1
+Version: 0.0.2
 Summary: Codon model package
 Author: CodonTeam
 Requires-Python: >=3.8
 License-File: LICENSE
 Requires-Dist: torch
+Requires-Dist: torchvision
 Requires-Dist: transformers
 Requires-Dist: tokenizers
 Requires-Dist: pandas
+Requires-Dist: numpy
 Requires-Dist: pyarrow
+Requires-Dist: pillow
 Dynamic: author
 Dynamic: license-file
 Dynamic: requires-dist

{codon_model-0.0.1 → codon_model-0.0.2}/codon/__init__.py

@@ -1,5 +1,5 @@
 from typing import Optional
 
-__version__ = '0.0.1'
+__version__ = '0.0.2'
 
 __seed__: Optional[int] = None

codon_model-0.0.2/codon/block/__init__.py

@@ -0,0 +1,83 @@
+from .attention import AttentionOutput, MultiHeadAttention
+from .codebook import LookupFreeQuantization, LookupFreeQuantizationOutput
+from .conv import (
+    CausalConv1d,
+    ConvBlock,
+    DepthwiseSeparableConv,
+    ResBasicBlock,
+    calculate_causal_layer,
+)
+from .embedding import (
+    BasicEmbedding,
+    BasicRotaryEmbedding,
+    InterleavedRotaryEmbedding,
+    RotaryEmbedding,
+    SinusoidalEmbedding,
+)
+from .film import FiLM, FiLMOutput
+from .fusion import (
+    CompactMultimodalPooling,
+    DiffusionMapsFusion,
+    GatedMultimodalUnit,
+    LowRankFusion,
+)
+from .lora import BasicLoRA, Conv1dLoRA, Conv2dLoRA, EmbeddingLoRA, LinearLoRA
+from .mlp import MLP
+from .moe import Expert, MoE, MoEInfo, MoEOutput
+from .pixelshuffle import PixelShuffleUpSample, UnPixelShuffleDownSample
+from .transformer import (
+    TransformerDecoderOutput,
+    TransformerDenseDecoder,
+    TransformerMoEDecoder,
+    _TransformerDecoder,
+)
+
+__all__ = [
+    # attention
+    'AttentionOutput',
+    'MultiHeadAttention',
+    # codebook
+    'LookupFreeQuantization',
+    'LookupFreeQuantizationOutput',
+    # conv
+    'CausalConv1d',
+    'ConvBlock',
+    'DepthwiseSeparableConv',
+    'ResBasicBlock',
+    'calculate_causal_layer',
+    # embedding
+    'BasicEmbedding',
+    'BasicRotaryEmbedding',
+    'InterleavedRotaryEmbedding',
+    'RotaryEmbedding',
+    'SinusoidalEmbedding',
+    # film
+    'FiLM',
+    'FiLMOutput',
+    # fusion
+    'CompactMultimodalPooling',
+    'DiffusionMapsFusion',
+    'GatedMultimodalUnit',
+    'LowRankFusion',
+    # lora
+    'BasicLoRA',
+    'Conv1dLoRA',
+    'Conv2dLoRA',
+    'EmbeddingLoRA',
+    'LinearLoRA',
+    # mlp
+    'MLP',
+    # moe
+    'Expert',
+    'MoE',
+    'MoEInfo',
+    'MoEOutput',
+    # pixelshuffle
+    'PixelShuffleUpSample',
+    'UnPixelShuffleDownSample',
+    # transformer
+    '_TransformerDecoder',
+    'TransformerDecoderOutput',
+    'TransformerDenseDecoder',
+    'TransformerMoEDecoder',
+]

codon_model-0.0.2/codon/block/attention.py

@@ -0,0 +1,173 @@
+from codon.base  import *
+from codon.block.embedding import BasicEmbedding
+from codon.ops.attention   import AttentionOutput, apply_attention
+
+
+class MultiHeadAttention(BasicModel):
+    ''' 
+    Multi-Head Attention module.
+    Supports Grouped Query Attention (GQA), QK Normalization, and Gating mechanisms.
+
+    Attributes:
+        q_proj (nn.Linear): Linear layer for query projection.
+        k_proj (nn.Linear): Linear layer for key projection.
+        v_proj (nn.Linear): Linear layer for value projection.
+        o_proj (nn.Linear): Linear layer for output projection.
+        q_norm (nn.RMSNorm, optional): Normalization layer for queries.
+        k_norm (nn.RMSNorm, optional): Normalization layer for keys.
+        g_proj (nn.Linear, optional): Linear layer for gating mechanism.
+    '''
+    def __init__(
+        self,
+        hidden_size,
+        num_heads,
+        num_kv_heads=None,
+        use_qk_norm=True,
+        use_gate=False,
+        dropout=0.1,
+        is_causal=True
+    ):
+        '''
+        Initialize the Multi-Head Attention module.
+
+        Args:
+            hidden_size (int): Size of the hidden layer.
+            num_heads (int): Number of attention heads.
+            num_kv_heads (int, optional): Number of key/value heads for GQA. 
+                                          If None, defaults to num_heads.
+            use_qk_norm (bool, optional): Whether to apply RMSNorm to queries and keys. 
+                                          Defaults to True.
+            use_gate (bool, optional): Whether to apply a gating mechanism. Defaults to False.
+            dropout (float, optional): Dropout probability. Defaults to 0.1.
+            is_causal (bool, optional): Whether to apply a causal mask. 
+                                        Defaults to True (for Decoder architectures).
+        '''
+        super(MultiHeadAttention, self).__init__()
+
+        if num_kv_heads is None: num_kv_heads = num_heads
+
+        assert hidden_size % num_heads == 0
+        assert num_heads % num_kv_heads == 0
+
+        self.hidden_size = hidden_size
+        self.num_heads  = num_heads
+        self.num_kv_heads = num_kv_heads
+        self.num_kv_queries = num_heads // num_kv_heads
+        self.head_dim  = hidden_size // num_heads
+        self.kv_dim = self.num_kv_heads * self.head_dim
+        self.use_qk_norm = use_qk_norm
+        self.use_gate = use_gate
+        self.dropout = dropout
+        self.is_causal = is_causal
+        
+        if use_qk_norm:
+            self.q_norm = nn.RMSNorm(self.head_dim)
+            self.k_norm = nn.RMSNorm(self.head_dim)
+        
+        if use_gate:
+            self.g_proj = nn.Linear(hidden_size, hidden_size)
+
+        self.q_proj = nn.Linear(hidden_size, hidden_size)
+        self.k_proj = nn.Linear(hidden_size, self.kv_dim)
+        self.v_proj = nn.Linear(hidden_size, self.kv_dim)
+        self.o_proj = nn.Linear(hidden_size, hidden_size)
+        
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        kv_states: torch.Tensor = None,
+        attention_mask: torch.Tensor = None,
+        output_attentions: bool = False,
+        position_emb: BasicEmbedding = None,
+        embedding_start: int = 0,
+        embedding_pos: torch.Tensor = None,
+        past_key_value: tuple[torch.Tensor, torch.Tensor] = None,
+        use_cache: bool = False
+    ) -> AttentionOutput:
+        ''' 
+        Perform forward pass of Multi-Head Attention.
+
+        Args:
+            hidden_states (torch.Tensor): Input hidden states.
+            kv_states (torch.Tensor, optional): Hidden states for keys/values. 
+                                                If None, uses hidden_states. Defaults to None.
+            attention_mask (torch.Tensor, optional): Attention mask. 
+                                                     Defaults to None.
+            output_attentions (bool, optional): Whether to output attention weights. 
+                                                Defaults to False.
+            position_emb (BasicEmbedding, optional): Positional embedding module. 
+                                                    Defaults to None.
+            embedding_start (int, optional): Starting position for embedding. Defaults to 0.
+            embedding_pos (torch.Tensor, optional): Explicit position indices for positional embedding. 
+                                                    Defaults to None.
+            past_key_value (tuple[torch.Tensor, torch.Tensor], optional): Past key-value cache. 
+                                                                          Defaults to None.
+            use_cache (bool, optional): Whether to use KV cache. Defaults to False.
+        
+        Returns:
+            AttentionOutput: Object containing output, attention weights, and KV cache.
+        '''
+        
+        if kv_states is None:
+            kv_states = hidden_states
+
+        batch_size, q_len, _ = hidden_states.shape
+        kv_len_input = kv_states.shape[1]
+
+        if self.use_gate:
+            G = torch.sigmoid(self.g_proj(hidden_states))
+        
+        Q = self.q_proj(hidden_states)
+        K = self.k_proj(kv_states)
+        V = self.v_proj(kv_states)
+
+        Q = Q.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        K = K.view(batch_size, kv_len_input, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        V = V.view(batch_size, kv_len_input, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        
+        if self.use_qk_norm:
+            Q = self.q_norm(Q)
+            K = self.k_norm(K)
+        
+        if position_emb is not None:
+            Q = position_emb(Q, start_pos=embedding_start, positions=embedding_pos)
+            K = position_emb(K, start_pos=embedding_start, positions=embedding_pos)
+        
+        current_key_value = None
+        if use_cache:
+            if past_key_value is not None:
+                past_k, past_v = past_key_value
+                K = torch.cat((past_k, K), dim=2)
+                V = torch.cat((past_v, V), dim=2)
+            current_key_value = (K, V)
+
+        kv_seq_len_total = K.shape[2]
+
+        if self.num_kv_queries > 1:
+            # [B, H_kv, 1, L, D] -> [B, H_kv, G, L, D]
+            K = K[:, :, None, :, :].expand(batch_size, self.num_kv_heads, self.num_kv_queries, kv_seq_len_total, self.head_dim)
+            V = V[:, :, None, :, :].expand(batch_size, self.num_kv_heads, self.num_kv_queries, kv_seq_len_total, self.head_dim)
+            
+            K = K.reshape(batch_size, self.num_heads, kv_seq_len_total, self.head_dim)
+            V = V.reshape(batch_size, self.num_heads, kv_seq_len_total, self.head_dim)
+            
+        attn_output = apply_attention(
+            Q, K, V, 
+            attention_mask=attention_mask, 
+            output_attentions=output_attentions,
+            is_causal=self.is_causal,
+            dropout=self.dropout if self.training else 0.0
+        )
+
+        output = attn_output.output
+        attention_weights = attn_output.attention_weights
+        output = output.transpose(1, 2).contiguous().view(batch_size, q_len, self.hidden_size)
+        output = self.o_proj(output)
+
+        if self.use_gate: output = output * G
+
+        return AttentionOutput(
+            output=output,
+            attention_weights=attention_weights,
+            past_key_value=current_key_value
+        )

codon_model-0.0.2/codon/block/bio/hebian.py

@@ -0,0 +1,223 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+
+from typing import Optional, Dict
+from dataclasses import dataclass
+
+from codon.base import BasicModel
+from codon.ops.bio import (
+    hebbian_update, 
+    oja_update, 
+    bcm_update, 
+    covariance_update, 
+    instar_update,
+    synaptic_scaling_update,
+    vogels_sprekeler_update,
+    reward_modulated_hebbian_update,
+    rate_based_stdp_update,
+    eligibility_trace_update
+)
+
+@dataclass
+class HebianOutput:
+    '''
+    Output of the Hebian layer.
+
+    Attributes:
+        output_tensor (torch.Tensor): The output/activation of the layer.
+        weight_updates (Dict[str, torch.Tensor]): A dictionary containing weight updates 
+                                                  for synapses.
+    '''
+    output_tensor: torch.Tensor
+    weight_updates: Dict[str, torch.Tensor]
+
+
+class Hebian(BasicModel):
+    '''
+    A layer implementing various biologically plausible Hebbian learning rules.
+
+    Attributes:
+        weight (nn.Parameter): Forward synaptic weights.
+        bias (nn.Parameter, optional): Bias term for the forward activation.
+    '''
+
+    def __init__(
+        self, 
+        in_features: int, 
+        out_features: int, 
+        learning_rate: float = 0.01, 
+        rule: str = 'oja', 
+        use_bias: bool = True,
+        auto_update: bool = False,
+        bcm_momentum: float = 0.1,
+        target_rate: float = 0.1,
+        trace_decay: float = 0.9,
+        activation: str = 'linear'
+    ) -> None:
+        '''
+        Initializes the Hebian layer.
+
+        Args:
+            in_features (int): Dimension of the input data.
+            out_features (int): Dimension of the output representation.
+            learning_rate (float, optional): Synaptic plasticity learning rate. Defaults to 0.01.
+            rule (str, optional): Learning rule ('hebbian', 'oja', 'bcm', 'covariance', 'instar', 
+                                  'scaling', 'vogels', 'reward_hebb', 'stdp', 'eligibility'). Defaults to 'oja'.
+            use_bias (bool, optional): Whether to use a bias term. Defaults to True.
+            auto_update (bool, optional): Automatically apply calculated weight updates in forward. Defaults to False.
+            bcm_momentum (float, optional): Momentum for BCM sliding threshold. Defaults to 0.1.
+            target_rate (float, optional): Desired average firing rate for homeostasis. Defaults to 0.1.
+            trace_decay (float, optional): Decay factor for eligibility trace. Defaults to 0.9.
+            activation (str, optional): Activation function ('linear', 'relu', 'sigmoid', 'tanh'). Defaults to 'linear'.
+        '''
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.learning_rate = learning_rate
+        self.rule = rule.lower()
+        self.use_bias = use_bias
+        self.auto_update = auto_update
+        self.bcm_momentum = bcm_momentum
+        self.target_rate = target_rate
+        self.trace_decay = trace_decay
+        self.activation = activation.lower()
+
+        self.weight = nn.Parameter(torch.Tensor(out_features, in_features), requires_grad=False)
+        
+        if self.use_bias:
+            self.bias = nn.Parameter(torch.Tensor(out_features), requires_grad=False)
+        else:
+            self.register_parameter('bias', None)
+            
+        if self.rule == 'bcm':
+            self.register_buffer('bcm_threshold', torch.zeros(out_features))
+        else:
+            self.bcm_threshold = None
+            
+        if self.rule == 'stdp':
+            self.register_buffer('prev_input', None)
+            self.register_buffer('prev_state', None)
+            
+        if self.rule == 'eligibility':
+            self.register_buffer('eligibility_trace', None)
+            
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        '''
+        Resets all synaptic parameters using Kaiming/Uniform initializations.
+        '''
+        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+        
+        if self.bias is not None:
+            fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight)
+            bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
+            nn.init.uniform_(self.bias, -bound, bound)
+
+    def _apply_activation(self, input_tensor: torch.Tensor) -> torch.Tensor:
+        '''
+        Applies the selected activation function.
+        
+        Args:
+            input_tensor (torch.Tensor): Input tensor.
+            
+        Returns:
+            torch.Tensor: Activated tensor.
+        '''
+        if self.activation == 'relu':
+            return F.relu(input_tensor)
+        elif self.activation == 'sigmoid':
+            return torch.sigmoid(input_tensor)
+        elif self.activation == 'tanh':
+            return torch.tanh(input_tensor)
+        return input_tensor
+
+    @torch.no_grad()
+    def forward(self, input_tensor: torch.Tensor, reward: Optional[torch.Tensor] = None) -> HebianOutput:
+        '''
+        Calculates the forward pass and biological synaptic updates.
+
+        Note:
+            Decorated with @torch.no_grad() to block global backpropagation, 
+            as learning is done via local plasticity rules.
+
+        Args:
+            input_tensor (torch.Tensor): The input data with shape (batch_size, in_features).
+            reward (Optional[torch.Tensor], optional): Global reward signal for 'reward_hebb' or 'eligibility' rules. Defaults to None.
+
+        Returns:
+            HebianOutput: Output containing the activation state and weight updates.
+        '''
+        r_state = F.linear(input_tensor, self.weight, self.bias)
+        r_state = self._apply_activation(r_state)
+        
+        updates: Dict[str, torch.Tensor] = {}
+        
+        if self.rule == 'hebbian':
+            updates['weight'] = hebbian_update(self.weight, input_tensor, r_state, self.learning_rate)
+        elif self.rule == 'oja':
+            updates['weight'] = oja_update(self.weight, input_tensor, r_state, self.learning_rate)
+        elif self.rule == 'bcm':
+            updates['weight'] = bcm_update(self.weight, input_tensor, r_state, self.bcm_threshold, self.learning_rate)
+            # Update sliding threshold: E[y^2]
+            current_y2 = torch.mean(r_state ** 2, dim=0)
+            self.bcm_threshold.mul_(1 - self.bcm_momentum).add_(current_y2, alpha=self.bcm_momentum)
+        elif self.rule == 'covariance':
+            updates['weight'] = covariance_update(self.weight, input_tensor, r_state, self.learning_rate)
+        elif self.rule == 'instar':
+            updates['weight'] = instar_update(self.weight, input_tensor, r_state, self.learning_rate)
+        elif self.rule == 'scaling':
+            updates['weight'] = synaptic_scaling_update(self.weight, r_state, target_rate=self.target_rate, learning_rate=self.learning_rate)
+        elif self.rule == 'vogels':
+            updates['weight'] = vogels_sprekeler_update(input_tensor, r_state, target_rate=self.target_rate, learning_rate=self.learning_rate)
+        elif self.rule == 'reward_hebb':
+            if reward is None:
+                raise ValueError("The 'reward_hebb' rule requires a reward signal to be passed to forward().")
+            updates['weight'] = reward_modulated_hebbian_update(input_tensor, r_state, reward, self.learning_rate)
+        elif self.rule == 'stdp':
+            if getattr(self, 'prev_input', None) is None or self.prev_input.shape != input_tensor.shape:
+                self.prev_input = input_tensor.clone().detach()
+                self.prev_state = r_state.clone().detach()
+                updates['weight'] = torch.zeros_like(self.weight)
+            else:
+                updates['weight'] = rate_based_stdp_update(input_tensor, self.prev_input, r_state, self.prev_state, self.learning_rate)
+                self.prev_input = input_tensor.clone().detach()
+                self.prev_state = r_state.clone().detach()
+        elif self.rule == 'eligibility':
+            current_hebbian = hebbian_update(self.weight, input_tensor, r_state, learning_rate=1.0)
+            
+            if getattr(self, 'eligibility_trace', None) is None or self.eligibility_trace.shape != self.weight.shape:
+                self.eligibility_trace = torch.zeros_like(self.weight)
+                
+            self.eligibility_trace = self.eligibility_trace * self.trace_decay + current_hebbian
+            
+            if reward is not None:
+                updates['weight'] = eligibility_trace_update(self.eligibility_trace, reward, self.learning_rate)
+            else:
+                updates['weight'] = torch.zeros_like(self.weight)
+        else:
+            raise ValueError(f'Unsupported learning rule: {self.rule}')
+
+        if self.auto_update:
+            self.apply_updates(updates)
+
+        return HebianOutput(
+            output_tensor=r_state,
+            weight_updates=updates
+        )
+
+    @torch.no_grad()
+    def apply_updates(self, updates: Dict[str, torch.Tensor]) -> None:
+        '''
+        Applies the calculated weight updates to the layer's parameters in-place.
+
+        Args:
+            updates (Dict[str, torch.Tensor]): A dictionary of parameter names and their updates.
+        '''
+        if 'weight' in updates and self.weight is not None:
+            self.weight.add_(updates['weight'])
+        
+        if 'bias' in updates and getattr(self, 'bias', None) is not None:
+            self.bias.add_(updates['bias'])