codon-model 0.0.4__tar.gz → 0.0.5__tar.gz

{codon_model-0.0.4/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.4
+Version: 0.0.5
 Summary: Codon model package
 Author: CodonTeam
 Requires-Python: >=3.8

{codon_model-0.0.4 → codon_model-0.0.5}/codon/__init__.py

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

{codon_model-0.0.4 → codon_model-0.0.5}/codon/base.py

@@ -5,6 +5,7 @@ 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
 
 
@@ -133,41 +134,106 @@ class BasicModel(nn.Module):
         
         return total
     
-    def load_pretrained(self, path: str) -> 'BasicModel':
+    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)
+            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) -> 'BasicModel':
+    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':

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

@@ -36,6 +36,7 @@ from .manifold import (
     BasicManifoldLinear, RiemannianManifoldLinear,
     BasicManifoldConv2d, RiemannianManifoldConv2d
 )
+from .fourier import MultiHeadFourier
 
 __all__ = [
     # attention
@@ -90,5 +91,7 @@ __all__ = [
     'BasicManifoldLinear',
     'RiemannianManifoldLinear',
     'BasicManifoldConv2d',
-    'RiemannianManifoldConv2d'
+    'RiemannianManifoldConv2d',
+    # Fourier
+    'MultiHeadFourier',
 ]

{codon_model-0.0.4 → 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.4 → 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.4 → codon_model-0.0.5}/codon/block/mlp.py

@@ -86,3 +86,41 @@ class MLP(BasicModel):
             x = self.dropout(x)
             x = self.fc2(x)
             return x
+    
+    @staticmethod
+    def SwiGLU(
+        in_features: int,
+        hidden_features: int = None,
+        out_features: int = None,
+        bias: bool = False,
+        dropout: float = 0.0,
+    ) -> 'MLP':
+        '''
+        Factory method to create a SwiGLU MLP module.
+        
+        SwiGLU formulation:
+            output = down_proj(SiLU(gate_proj(x)) * up_proj(x))
+            
+        Args:
+            in_features (int): Input dimension.
+            hidden_features (int): Intermediate dimension (gate & up proj).
+            out_features (int, optional): Output dimension. Defaults to in_features.
+            bias (bool, optional): Whether to use bias. LLMs typically set False. Defaults to False.
+            dropout (float, optional): Dropout rate. Usually 0.0 for SwiGLU. Defaults to 0.0.
+            
+        Returns:
+            MLP: Configured SwiGLU module.
+        '''
+        if hidden_features is None:
+            h = int(in_features * 8 / 3)
+            hidden_features = (h + 127) // 128 * 128
+            
+        return MLP(
+            in_features=in_features,
+            hidden_features=hidden_features,
+            out_features=out_features,
+            bias=bias,
+            use_gate=True,
+            dropout=dropout,
+            act_layer='silu'
+        )

{codon_model-0.0.4 → codon_model-0.0.5}/codon/block/transformer.py

@@ -5,6 +5,7 @@ from typing import Union, Optional, Tuple
 from codon.utils.safecode  import safecode
 from codon.block.embedding import BasicEmbedding
 from codon.block.attention import MultiHeadAttention, AttentionOutput
+from codon.block.fourier   import MultiHeadFourier
 from codon.block.mlp       import MLP
 from codon.block.moe       import MoE, MoEOutput
 
@@ -75,6 +76,8 @@ class _TransformerDecoder(BasicModel):
         use_qk_norm=True,
         use_attn_gate=False,
         dropout=0.1,
+        attn_type: str = 'multihead',
+        attn_bias: bool=False,
         idx: Union[int, str] = None
     ):
         '''
@@ -97,17 +100,32 @@ class _TransformerDecoder(BasicModel):
         self.num_kv_heads = num_kv_heads
         self.use_qk_norm = use_qk_norm
         self.use_attn_gate = use_attn_gate
+        self.attn_bias = attn_bias
+        self.attn_type = attn_type.lower()
 
         self.attn_norm = nn.RMSNorm(model_dim)
-        self.attn = MultiHeadAttention(
-            hidden_size=model_dim,
-            num_heads=num_heads,
-            num_kv_heads=num_kv_heads,
-            use_qk_norm=use_qk_norm,
-            use_gate=use_attn_gate,
-            dropout=dropout,
-            is_causal=True
-        )
+        
+        if self.attn_type in ['multihead', 'mha']:
+            self.attn = MultiHeadAttention(
+                hidden_size=model_dim,
+                num_heads=num_heads,
+                num_kv_heads=num_kv_heads,
+                use_qk_norm=use_qk_norm,
+                use_gate=use_attn_gate,
+                dropout=dropout,
+                is_causal=True,
+                bias=attn_bias
+            )
+
+        elif self.attn_type in ['fourier', 'mhf']:
+            self.attn = MultiHeadFourier(
+                hidden_size=model_dim,
+                num_heads=num_heads
+            )
+
+        else:
+            raise ValueError(f"Unsupported attn_type: {attn_type}. Choose from 'multihead', 'linear', or 'gated_linear'.")
+
         self.fn_norm = nn.RMSNorm(model_dim)
         self.dropout = nn.Dropout(dropout)
     
@@ -228,7 +246,11 @@ class TransformerDenseDecoder(_TransformerDecoder):
         use_mlp_gate: bool = False,
         use_qk_norm: bool = True,
         use_attn_gate: bool = False,
+        attn_type: str = 'multihead',
+        use_swiglu: bool = False,
         dropout: float = 0.1,
+        attn_bias: bool=False,
+        mlp_bias: bool=False,
         idx: Union[int, str] = None
     ):
         '''
@@ -252,15 +274,18 @@ class TransformerDenseDecoder(_TransformerDecoder):
             use_qk_norm=use_qk_norm,
             use_attn_gate=use_attn_gate,
             dropout=dropout,
-            idx=idx
+            attn_type=attn_type,
+            idx=idx,
+            attn_bias=attn_bias
         )
         self.mlp = MLP(
             in_features=model_dim,
             hidden_features=int(model_dim * mlp_ratio),
             out_features=model_dim,
             use_gate=use_mlp_gate,
-            dropout=dropout
-        )
+            dropout=dropout,
+            bias=mlp_bias
+        ) if not use_swiglu else MLP.SwiGLU(model_dim)
     
     def flow(self, x: torch.Tensor) -> FlowOutput:
         '''
@@ -294,7 +319,9 @@ class TransformerMoEDecoder(_TransformerDecoder):
         use_expert_gate: bool = False,
         use_qk_norm: bool = True,
         use_attn_gate: bool = False,
+        attn_type: str = 'multihead',
         dropout: float = 0.1,
+        attn_bias: bool=False,
         idx: Union[int, str] = None
     ):
         '''
@@ -321,7 +348,9 @@ class TransformerMoEDecoder(_TransformerDecoder):
             use_qk_norm=use_qk_norm,
             use_attn_gate=use_attn_gate,
             dropout=dropout,
-            idx=idx
+            attn_type=attn_type,
+            idx=idx,
+            attn_bias=attn_bias
         )
         self.moe = MoE(
             model_dim=model_dim,

codon_model-0.0.5/codon/cli.py

@@ -0,0 +1,48 @@
+import argparse
+import os
+
+from codon.dev.hash import hash_target
+from codon.dev.clear import clear_pycache
+
+def handle_hash(args):
+    '''处理 codon hash 命令'''
+    try:
+        results = hash_target(args.filename_or_dirname)
+        for path, md5_val in results.items():
+            print(f'{md5_val}  {path}')
+    except ValueError as e:
+        print(e)
+    except Exception as e:
+        print(f'发生意外错误: {e}')
+
+def handle_clear(args):
+    '''处理 codon clear 命令'''
+    if not os.path.isdir(args.target_path):
+        print(f'错误: 目录不存在 - {args.target_path}')
+        return
+    clear_pycache(args.target_path)
+
+def main():
+    parser = argparse.ArgumentParser(
+        prog='codon',
+        description='Codon Model 统一命令行工具箱'
+    )
+    
+    subparsers = parser.add_subparsers(title='可用子命令', dest='command', required=True)
+
+    # codon hash
+    parser_hash = subparsers.add_parser('hash', help='计算文件或目录的 MD5 值')
+    parser_hash.add_argument('filename_or_dirname', type=str, help='需要计算 MD5 的文件或文件夹路径')
+    parser_hash.set_defaults(func=handle_hash)
+
+    # codon clear
+    parser_clear = subparsers.add_parser('clear', help='递归清理目录下的所有 __pycache__ 缓存')
+    parser_clear.add_argument('target_path', type=str, nargs='?', default='.', 
+                              help='要清理的根目录路径 (默认为当前所在目录)')
+    parser_clear.set_defaults(func=handle_clear)
+
+    args = parser.parse_args()
+    args.func(args)
+
+if __name__ == '__main__':
+    main()

codon_model-0.0.5/codon/dev/clear.py

@@ -0,0 +1,41 @@
+import os
+import shutil
+import argparse
+
+def clear_pycache(target_dir: str):
+    '''
+    递归遍历并删除目标路径下的所有 __pycache__ 文件夹。
+    
+    参数:
+        target_dir (str): 需要清理的根目录路径。
+    '''
+    deleted_count = 0
+    target_dir = os.path.abspath(target_dir)
+
+    print(f'开始清理 {target_dir} 下的 __pycache__ ...')
+    
+    for root, dirs, files in os.walk(target_dir):
+        if '__pycache__' in dirs:
+            cache_dir = os.path.join(root, '__pycache__')
+            try:
+                shutil.rmtree(cache_dir)
+                print(f'Deleted: {cache_dir}')
+                deleted_count += 1
+                
+                dirs.remove('__pycache__')
+            except Exception as e:
+                print(f'Error deleting {cache_dir}: {e}')
+                
+    print(f'清理完成。共移除了 {deleted_count} 个 __pycache__ 缓存目录。')
+
+def main():
+    parser = argparse.ArgumentParser(description='递归清理目标路径下的所有 __pycache__ 文件夹')
+    parser.add_argument('target_path', type=str, nargs='?', default='.', 
+                        help='要清理的根目录路径 (默认为当前所在目录)')
+    args = parser.parse_args()
+
+    if not os.path.isdir(args.target_path):
+        print(f'错误: 目录不存在 - {args.target_path}')
+        return
+
+    clear_pycache(args.target_path)

codon_model-0.0.5/codon/dev/hash.py

@@ -0,0 +1,54 @@
+import hashlib
+import os
+import argparse
+from typing import Dict
+
+from tqdm import tqdm
+
+def calculate_md5(file_path: str, chunk_size: int = 8192) -> str:
+    '''
+    计算单个文件的 MD5 摘要，带有阅后即焚的进度条。
+    '''
+    md5_hash = hashlib.md5()
+    try:
+        total_size = os.path.getsize(file_path)
+        with tqdm(total=total_size, unit='B', unit_scale=True, unit_divisor=1024, 
+                  desc=f'Hashing {os.path.basename(file_path)}', leave=False) as pbar:
+            with open(file_path, 'rb') as f:
+                for chunk in iter(lambda: f.read(chunk_size), b''):
+                    md5_hash.update(chunk)
+                    pbar.update(len(chunk))
+                    
+        return md5_hash.hexdigest()
+    except Exception as e:
+        return f'Error: {e}'
+
+def hash_target(target_path: str) -> Dict[str, str]:
+    results = {}
+    if os.path.isfile(target_path):
+        results[target_path] = calculate_md5(target_path)
+    elif os.path.isdir(target_path):
+        for root, _, files in os.walk(target_path):
+            for file in files:
+                file_path = os.path.join(root, file)
+                results[file_path] = calculate_md5(file_path)
+    else:
+        raise ValueError(f'无效的路径: {target_path}')
+    
+    return results
+
+def main():
+    parser = argparse.ArgumentParser(description='计算文件或目录下所有文件的 MD5 值')
+    parser.add_argument('target', type=str, help='需要计算 MD5 的文件或文件夹路径')
+    args = parser.parse_args()
+
+    try:
+        results = hash_target(args.target)
+        for path, md5_val in results.items():
+            path = path.replace('\\', '/')
+            print(f'{md5_val} ➜ {path}')
+    except ValueError as e:
+        print(e)
+
+if __name__ == '__main__':
+    main()

{codon_model-0.0.4 → codon_model-0.0.5}/codon/exp/block/bio.py

@@ -3,7 +3,7 @@ from codon.base import *
 from typing import List, Optional
 from dataclasses import dataclass
 
-from codon.utils.dataset.base import CodonDataset
+from codon.utils.data.base import CodonDataset
 
 @dataclass
 class EABStats:

{codon_model-0.0.4 → codon_model-0.0.5}/codon/kit/train/__init__.py

@@ -1,6 +1,8 @@
 from .vision import auto_vision_train, AutoVisionTrainResult
+from .language import run_sanity_check
 
 __all__ = [
     'auto_vision_train',
-    'AutoVisionTrainResult'
+    'AutoVisionTrainResult',
+    'run_sanity_check'
 ]