codon-model 0.0.3b2__tar.gz → 0.0.5__tar.gz

{codon_model-0.0.3b2/codon_model.egg-info → codon_model-0.0.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: codon-model
-Version: 0.0.3b2
+Version: 0.0.5
 Summary: Codon model package
 Author: CodonTeam
 Requires-Python: >=3.8

{codon_model-0.0.3b2 → codon_model-0.0.5}/codon/__init__.py

@@ -1,5 +1,5 @@
 from typing import Optional
 
-__version__ = '0.0.3b2'
+__version__ = '0.0.5'
 
 __seed__: Optional[int] = None

{codon_model-0.0.3b2 → codon_model-0.0.5}/codon/base.py

@@ -1,9 +1,11 @@
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
 
 from typing import Callable, Any, Iterator, Union
 
 from safetensors.torch import save_model as safe_save_model
+from safetensors.torch import save_file  as safe_save_file
 from safetensors.torch import load_model as safe_load_model
 
 
@@ -132,36 +134,126 @@ class BasicModel(nn.Module):
         
         return total
     
-    def load_pretrained(self, path: str) -> None:
+    def load_pretrained(self, path: str, strict: bool = False) -> 'BasicModel':
         '''
         Load a pretrained model from a file.
-
         Args:
             path (str): The path to the model file.
+            strict (bool, optional): Whether to strictly enforce that the keys
+                                     in state_dict match. Defaults to False.
         '''
         if path.endswith('.safetensors'):
-            safe_load_model(self, path)
-            return
-
+            safe_load_model(self, path, strict=strict)
+            return self
         state_dict = torch.load(path, map_location=self.device)
-
         if isinstance(state_dict, dict):
             if 'model_state_dict' in state_dict:
                 state_dict = state_dict['model_state_dict']
             elif 'state_dict' in state_dict:
                 state_dict = state_dict['state_dict']
+            elif 'model' in state_dict:
+                state_dict = state_dict['model']
         
-        self.load_state_dict(state_dict)
+        clean_state_dict = {k.replace('_orig_mod.', ''): v for k, v in state_dict.items()}
+        
+        self.load_state_dict(clean_state_dict, strict=strict)
+        return self
     
-    def save_pretrained(self, path: str) -> None:
+    def save_pretrained(
+            self, 
+            path: str, 
+            trainable_only: bool = False, 
+            include_buffer: bool = True, 
+            exclude_modules: list[Union[type, nn.Module]] = None,
+            only: list[str] = None,
+            exclude: list[str] = None
+        ) -> 'BasicModel':
         '''
         Save the model to a file.
 
         Args:
             path (str): The path to save the model file.
+            trainable_only (bool, optional): If True, only save parameters that require gradients.
+            include_buffer (bool, optional): If False, exclude registered buffers from the saved file.
+            exclude_modules (list[Union[type, nn.Module]], optional): Module types or instances to exclude.
+            only (list[str], optional): If provided, only save parameters whose keys contain ANY of these strings.
+            exclude (list[str], optional): If provided, exclude parameters whose keys contain ANY of these strings.
         '''
+        state_dict = self.state_dict()
+        is_modified = False
+        
+        exclude_prefixes = []
+        if exclude_modules:
+            exclude_types = tuple(t for t in exclude_modules if isinstance(t, type))
+            exclude_instances = set(m for m in exclude_modules if not isinstance(m, type))
+            
+            for name, module in self.named_modules():
+                if module in exclude_instances or (exclude_types and isinstance(module, exclude_types)):
+                    if name != '': exclude_prefixes.append(name + '.')
+        exclude_prefixes = tuple(exclude_prefixes)
+
+        has_filter = trainable_only or not include_buffer or exclude_prefixes or only or exclude
+
+        if has_filter:
+            trainable_names = {name for name, p in self.named_parameters() if p.requires_grad}
+            buffer_names = {name for name, _ in self.named_buffers()}
+            
+            filtered_dict = {}
+            for key, tensor in state_dict.items():
+                keep = True
+                
+                if exclude_prefixes and key.startswith(exclude_prefixes):
+                    keep = False
+                
+                elif exclude and any(kw in key for kw in exclude):
+                    keep = False
+                
+                elif only and not any(kw in key for kw in only):
+                    keep = False
+                
+                else:
+                    is_buffer = key in buffer_names
+                    if not include_buffer and is_buffer:
+                        keep = False
+                    elif trainable_only and not is_buffer and key not in trainable_names:
+                        keep = False
+                
+                if keep:
+                    filtered_dict[key] = tensor
+                else:
+                    is_modified = True
+            
+            if is_modified:
+                state_dict = filtered_dict
+
         if path.endswith('.safetensors'):
-            safe_save_model(self, path)
+            if not is_modified:
+                safe_save_model(self, path)
+            else:
+                safe_save_file(state_dict, path)
         else:
-            state_dict = self.state_dict()
             torch.save(state_dict, path)
+            
+        return self
+    
+    def freeze(self) -> 'BasicModel':
+        '''
+        Freeze all parameters in the model by setting requires_grad to False.
+
+        Returns:
+            BasicModel: The model itself for method chaining.
+        '''
+        for param in self.parameters():
+            param.requires_grad = False
+        return self
+
+    def unfreeze(self) -> 'BasicModel':
+        '''
+        Unfreeze all parameters in the model by setting requires_grad to True.
+
+        Returns:
+            BasicModel: The model itself for method chaining.
+        '''
+        for param in self.parameters():
+            param.requires_grad = True
+        return self

{codon_model-0.0.3b2 → codon_model-0.0.5}/codon/block/__init__.py

@@ -31,6 +31,12 @@ from .transformer import (
     TransformerMoEDecoder,
     _TransformerDecoder,
 )
+from .manifold import (
+    MainfoldLoss,
+    BasicManifoldLinear, RiemannianManifoldLinear,
+    BasicManifoldConv2d, RiemannianManifoldConv2d
+)
+from .fourier import MultiHeadFourier
 
 __all__ = [
     # attention
@@ -80,4 +86,12 @@ __all__ = [
     'TransformerDecoderOutput',
     'TransformerDenseDecoder',
     'TransformerMoEDecoder',
+    # manifold
+    'MainfoldLoss',
+    'BasicManifoldLinear',
+    'RiemannianManifoldLinear',
+    'BasicManifoldConv2d',
+    'RiemannianManifoldConv2d',
+    # Fourier
+    'MultiHeadFourier',
 ]

{codon_model-0.0.3b2 → codon_model-0.0.5}/codon/block/attention.py

@@ -1,6 +1,9 @@
 from codon.base  import *
 from codon.block.embedding import BasicEmbedding
-from codon.ops.attention   import AttentionOutput, apply_attention
+from codon.ops.attention   import (
+    AttentionOutput,
+    apply_attention
+)
 
 
 class MultiHeadAttention(BasicModel):
@@ -25,6 +28,7 @@ class MultiHeadAttention(BasicModel):
         use_qk_norm=True,
         use_gate=False,
         dropout=0.1,
+        bias: bool=True,
         is_causal=True
     ):
         '''
@@ -65,12 +69,12 @@ class MultiHeadAttention(BasicModel):
             self.k_norm = nn.RMSNorm(self.head_dim)
         
         if use_gate:
-            self.g_proj = nn.Linear(hidden_size, hidden_size)
+            self.g_proj = nn.Linear(hidden_size, hidden_size, bias=bias)
 
-        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)
+        self.q_proj = nn.Linear(hidden_size, hidden_size, bias=bias)
+        self.k_proj = nn.Linear(hidden_size, self.kv_dim, bias=bias)
+        self.v_proj = nn.Linear(hidden_size, self.kv_dim, bias=bias)
+        self.o_proj = nn.Linear(hidden_size, hidden_size, bias=bias)
         
     def forward(
         self,
@@ -171,3 +175,4 @@ class MultiHeadAttention(BasicModel):
             attention_weights=attention_weights,
             past_key_value=current_key_value
         )
+

codon_model-0.0.5/codon/block/bio/__init__.py

@@ -0,0 +1,9 @@
+from .hebian import HebianOutput, Hebian
+from .predictive import PredictiveCodingOutput, PredictiveCoding
+
+__all__ = [
+    'HebianOutput',
+    'Hebian',
+    'PredictiveCodingOutput',
+    'PredictiveCoding'
+]

{codon_model-0.0.3b2 → codon_model-0.0.5}/codon/block/embedding.py

@@ -69,7 +69,7 @@ class SinusoidalEmbedding(BasicEmbedding):
         pe[:, 1::2] = torch.cos(position * div_term)
         
         pe = pe.unsqueeze(0)
-        self.register_buffer('pe', pe)
+        self.register_buffer('pe', pe, persistent=False)
 
     def forward(self, x: torch.Tensor, positions: torch.Tensor = None, start_pos: int = 0) -> torch.Tensor:
         '''
@@ -132,8 +132,8 @@ class BasicRotaryEmbedding(BasicEmbedding):
 
         emb = torch.cat((freqs, freqs), dim=-1)
 
-        self.register_buffer('cos_cached', emb.cos())
-        self.register_buffer('sin_cached', emb.sin())
+        self.register_buffer('cos_cached', emb.cos(), persistent=False)
+        self.register_buffer('sin_cached', emb.sin(), persistent=False)
 
     def _rotate_half(self, x: torch.Tensor) -> torch.Tensor:
         '''

codon_model-0.0.5/codon/block/fourier.py

@@ -0,0 +1,106 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from codon.base import BasicModel
+from codon.ops.attention import AttentionOutput
+from codon.ops.fourier   import apply_fourier_mixing
+
+
+class MultiHeadFourier(BasicModel):
+    '''
+    Multi-Head Fourier (MHF) module from 'Caracal: Causal Architecture via Spectral Mixing'.
+    Replaces dense Attention with O(L log L) frequency-domain mixing.
+    '''
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        **kwargs
+    ):
+        super().__init__()
+        
+        assert hidden_size % num_heads == 0, 'hidden_size must be divisible by num_heads'
+        
+        self.d_model = hidden_size
+        self.n_heads = num_heads
+        self.d_head = hidden_size // num_heads
+        
+        self.pre_conv = nn.Conv1d(
+            in_channels=self.d_model, 
+            out_channels=self.d_model,
+            kernel_size=3, 
+            groups=self.d_model, 
+            bias=False
+        )
+        self.ln = nn.LayerNorm(self.d_model)
+        self.silu = nn.SiLU()
+        
+        self.W_V = nn.Linear(self.d_model, self.d_model)
+        self.W_G1 = nn.Linear(self.d_model, self.d_model)
+        self.W_G2 = nn.Conv1d(
+            in_channels=self.d_model, 
+            out_channels=self.d_model,
+            kernel_size=1, 
+            groups=self.n_heads
+        )
+        
+        self.linear = nn.Linear(self.d_model, self.d_model)
+        self.linear.NEED_SCALE_INIT = 1
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        use_cache: bool = False,
+        past_key_value: tuple = None,
+        **kwargs
+    ) -> AttentionOutput:
+        
+        batch_size, seq_len, _ = hidden_states.size()
+        
+        x_permuted = hidden_states.permute(0, 2, 1)  # [B, D, L]
+        
+        if past_key_value is not None:
+            conv_state, v_cache, g_cache = past_key_value
+            x_padded = torch.cat([conv_state, x_permuted], dim=2)
+            x = self.pre_conv(x_padded).permute(0, 2, 1)  # [B, 1, D]
+            
+            new_conv_state = x_padded[:, :, 1:]
+        else:
+            padding = self.pre_conv.kernel_size[0] - 1
+            x_padded = F.pad(x_permuted, (padding, 0))
+            x = self.pre_conv(x_padded).permute(0, 2, 1)  # [B, L, D]
+            
+            new_conv_state = x_padded[:, :, -padding:] if use_cache else None
+
+        x_norm = self.ln(x)
+        
+        x_v = self.W_V(x_norm).reshape(batch_size, seq_len, self.n_heads, self.d_head).transpose(1, 2)
+        
+        x_g = self.W_G1(x_norm).transpose(1, 2)
+        x_g = self.W_G2(self.silu(x_g)).transpose(1, 2)
+        x_g = x_g.reshape(batch_size, seq_len, self.n_heads, self.d_head).transpose(1, 2)
+        
+        if past_key_value is not None:
+            new_v_cache = torch.cat([v_cache, x_v], dim=2)
+            new_g_cache = torch.cat([g_cache, x_g], dim=2)
+            
+            g_flipped = torch.flip(new_g_cache, dims=[2])
+            
+            x_mixed = torch.sum(new_v_cache * g_flipped, dim=2, keepdim=True)
+        else:
+            x_mixed = apply_fourier_mixing(x_v, x_g, seq_len)
+            
+            new_v_cache = x_v if use_cache else None
+            new_g_cache = x_g if use_cache else None
+
+        x_mixed = x_mixed.transpose(1, 2).contiguous().reshape(batch_size, seq_len, self.d_model)
+        output = self.linear(x_mixed)
+        
+        current_key_value = (new_conv_state, new_v_cache, new_g_cache) if use_cache else None
+        
+        return AttentionOutput(
+            output=output,
+            attention_weights=None,
+            past_key_value=current_key_value
+        )

{codon_model-0.0.3b2 → codon_model-0.0.5}/codon/block/lora.py

@@ -50,6 +50,8 @@ class BasicLoRA(BasicModel):
         super().__init__()
         self.gradient_checkpointing = gradient_checkpointing
         self.original_layer = original_layer
+        self.lora_dropout_p = lora_dropout
+        self.merge_weights = merge_weights
 
         # Freeze original layer
         for p in self.original_layer.parameters():