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

{codon_model-0.0.4/codon_model.egg-info → codon_model-0.0.5a2}/PKG-INFO

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

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

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

{codon_model-0.0.4 → codon_model-0.0.5a2}/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,102 @@ 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']
         
-        self.load_state_dict(state_dict)
-
+        self.load_state_dict(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.5a2}/codon/motif/base.py

@@ -1,12 +1,8 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
+from codon.base import *
 
-from typing import Callable, Any, Iterator, Union, Optional, List, Tuple
+from typing import Optional, List, Tuple
 from dataclasses import dataclass
 
-from codon.base import BasicModel
-
 
 @dataclass
 class AutoVisionEncoderOutput:
@@ -77,7 +73,7 @@ class CausalLanguageModel(BasicModel):
         self,
         input_ids: torch.Tensor,
         max_new_tokens: int = 100,
-        temperature: float = 1.0,
+        temperature: float = 0.7,
         top_k: int = None,
         eos_token_id: int = None
     ) -> torch.Tensor:
@@ -88,7 +84,7 @@ class CausalLanguageModel(BasicModel):
             input_ids (torch.Tensor): Input token IDs with shape [batch, seq_len].
             max_new_tokens (int): Maximum number of new tokens to generate. Defaults to 100.
             temperature (float): Sampling temperature. Higher values increase randomness.
-                                 Defaults to 1.0.
+                                 Defaults to 0.7.
             top_k (int, optional): If set, sample only from top k tokens. Defaults to None.
             eos_token_id (int, optional): End-of-sequence token ID. If None, generation
                                           stops after max_new_tokens. Defaults to None.

codon_model-0.0.5a2/codon/motif/data/__init__.py

@@ -0,0 +1,5 @@
+from .prev1 import MotifPrev1
+
+__all__ = [
+    'MotifPrev1'
+]

codon_model-0.0.5a2/codon/motif/data/prev1.py

@@ -0,0 +1,85 @@
+from codon.utils.dataset.base import CodonDataset
+from transformers import PreTrainedTokenizerFast
+from codon.utils.tokens import PackedTokenizer
+
+import os
+import bisect
+import pyarrow.parquet as pq
+from pathlib import Path
+from typing  import Any, Optional, Dict, Union
+from tqdm    import tqdm
+
+class MotifPrev1(CodonDataset):
+    '''
+    A concrete implementation of CodonDataset for loading Motif data from a directory.
+    Optimized with O(1) lazy loading via Parquet metadata or full memory mapping initialization.
+    '''
+
+    def __init__(self, path: str, mode: str = 'lazy') -> None:
+        if not os.path.isdir(path):
+            raise NotADirectoryError(f'{path}')
+
+        if mode not in ['lazy', 'full']:
+            raise ValueError('')
+
+        self.path = path
+        self.mode = mode
+        self.tokenizer: Optional[PreTrainedTokenizerFast] = None
+        
+        self.file_paths = sorted(list(Path(path).glob('*.parquet')))
+        if not self.file_paths:
+            raise FileNotFoundError(f'{path}')
+
+        self.cum_sizes = []
+        self._table_cache: Dict[int, Any] = {}
+        current_total = 0
+        
+        with tqdm(total=len(self.file_paths), desc=f'Loading Dataset ({mode})', leave=False) as pbar:
+            for idx, fp in enumerate(self.file_paths):
+                if mode == 'lazy':
+                    meta = pq.read_metadata(fp)
+                    current_total += meta.num_rows
+                elif mode == 'full':
+                    table = pq.read_table(fp, memory_map=True)
+                    self._table_cache[idx] = table
+                    current_total += table.num_rows
+                
+                self.cum_sizes.append(current_total)
+                pbar.update(1)
+            
+        self.total_rows = current_total
+
+    def set_tokenizer(self, tokenizer: Union[PreTrainedTokenizerFast, PackedTokenizer]) -> 'MotifPrev1':
+        if isinstance(tokenizer, PackedTokenizer):
+            tokenizer = tokenizer.fast_tokenizer
+        self.tokenizer = tokenizer
+        return self
+
+    def __len__(self) -> int:
+        return self.total_rows
+
+    def __getitem__(self, idx: int) -> Dict[str, Any]:
+        if idx < 0 or idx >= self.total_rows:
+            raise IndexError(f'Index {idx} out of bounds for dataset length {self.total_rows}.')
+
+        file_idx = bisect.bisect_right(self.cum_sizes, idx)
+        local_idx = idx if file_idx == 0 else idx - self.cum_sizes[file_idx - 1]
+        
+        if file_idx not in self._table_cache:
+            self._table_cache[file_idx] = pq.read_table(self.file_paths[file_idx], memory_map=True)
+            
+        table = self._table_cache[file_idx]
+        
+        content_str = table.column('content')[local_idx].as_py()
+        
+        record: Dict[str, Any] = {
+            'content': content_str
+        }
+        
+        if 'tag' in table.column_names:
+            record['tag'] = table.column('tag')[local_idx].as_py()
+            
+        if self.tokenizer is not None:
+            record['input_ids'] = self.tokenizer.encode(content_str)
+            
+        return record

{codon_model-0.0.4 → codon_model-0.0.5a2}/codon/motif/motif_a1.py

@@ -34,9 +34,9 @@ class MotifA1(CausalLanguageModel):
                 top_k=1,
                 num_experts=3,
                 num_shared_experts=1,
-                use_expert_gate=True,
+                use_expert_gate=False,
                 use_qk_norm=True,
-                use_attn_gate=True,
+                use_attn_gate=False,
                 dropout=dropout,
                 idx=str(idx)
             )

{codon_model-0.0.4 → codon_model-0.0.5a2}/codon/ops/attention.py

@@ -5,7 +5,6 @@ import math
 from dataclasses import dataclass
 from typing      import Optional, Tuple
 
-from torch.nn.attention import SDPBackend, sdpa_kernel
 
 @dataclass
 class AttentionOutput:
@@ -77,21 +76,16 @@ def apply_attention(
             is_causal = True
         
         try:
-            with sdpa_kernel([
-                SDPBackend.FLASH_ATTENTION,
-                SDPBackend.CUDNN_ATTENTION
-            ]):
-                output = F.scaled_dot_product_attention(
-                    query_states, 
-                    key_states, 
-                    value_states, 
-                    attn_mask=attention_mask,
-                    is_causal=is_causal,
-                    dropout_p=dropout
-                )
+            output = F.scaled_dot_product_attention(
+                query_states, 
+                key_states, 
+                value_states, 
+                attn_mask=attention_mask,
+                is_causal=is_causal,
+                dropout_p=dropout
+            )
             return AttentionOutput(output=output, attention_weights=None)
-        except RuntimeError: 
-            pass
+        except RuntimeError: pass
     # Manual Fallback Path
     d_k = query_states.size(-1)
     scores = torch.matmul(query_states, key_states.transpose(-2, -1)) / math.sqrt(d_k)

{codon_model-0.0.4 → codon_model-0.0.5a2}/codon/ops/pixelshuffle.py

@@ -1,5 +1,4 @@
 import torch
-from typing import Tuple
 
 
 def pixel_shuffle(

{codon_model-0.0.4 → codon_model-0.0.5a2}/codon/utils/tokens.py

@@ -37,10 +37,10 @@ class TokenizerTrainerResult:
 core_tokens = ['[pad]', '[unk]', '[sep]', '[cls]']
 chat_tokens = [
     '[im_start]', '[im_end]',
-    '[system]', '[user]', '[model]', '[tool]', '[train]',
+    '[system]', '[user]', '[model]', '[tool]',
     '[interruption]', '[fim]',
 ]
-reasoning_tokens = ['[cot_start]', '[cot_end]', '[verification]', '[solution]']
+reasoning_tokens = ['[cot_start]', '[cot_end]']
 code_tokens = ['[fim_pre]', '[fim_mid]', '[fim_suf]']
 tool_tokens = ['[tool_start]', '[tool_name]', '[tool_args]', '[tool_end]']
 
@@ -59,7 +59,6 @@ base_special_tokens = (
 )
 
 base_special_tokens += [f'[unused_{i}]' for i in range(len(base_special_tokens), 64)]
-base_special_tokens += [f'[mask_{i}]' for i in range(32)]
 
 chat_template = (
     "{% for message in messages %}"
@@ -83,8 +82,6 @@ chat_template = (
                 "{{ '[model]' }}"
             "{% elif message['role'] == 'tool' %}"
                 "{{ '[tool]' }}"
-            "{% elif message['role'] == 'train' %}"
-                "{{ '[train]' }}"
             "{% else %}"
                 "{{ message['role'] }}"
             "{% endif %}"

{codon_model-0.0.4 → codon_model-0.0.5a2/codon_model.egg-info}/PKG-INFO

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

{codon_model-0.0.4 → codon_model-0.0.5a2}/codon_model.egg-info/SOURCES.txt

@@ -35,6 +35,8 @@ codon/motif/__init__.py
 codon/motif/base.py
 codon/motif/motif_a1.py
 codon/motif/motif_v1.py
+codon/motif/data/__init__.py
+codon/motif/data/prev1.py
 codon/ops/__init__.py
 codon/ops/attention.py
 codon/ops/bio.py