codon-model 0.0.1__py3-none-any.whl

codon/model/patch_disc.py

@@ -0,0 +1,151 @@
+import math
+
+from codon.base import *
+from codon.block.conv import ConvBlock
+
+
+class PatchDiscriminator(BasicModel):
+    '''
+    PatchGAN discriminator.
+
+    The output is not a scalar, but an N x N matrix, where each point represents
+    whether the corresponding patch is real or fake.
+
+    Attributes:
+        main (nn.Sequential): The main sequential model.
+    '''
+
+    def __init__(
+        self,
+        in_channels: int = 3,
+        hidden_dim: int = 64,
+        num_layers: int = 3,
+        norm: str = 'batch',
+        activation: str = 'leaky_relu',
+        leaky_relu: float = 0.2
+    ) -> None:
+        '''
+        Initialize the PatchDiscriminator.
+
+        Args:
+            in_channels (int): Number of input channels. Defaults to 3.
+            hidden_dim (int): Base number of filters (channels) in the discriminator. Defaults to 64.
+            num_layers (int): Number of layers in the discriminator. Defaults to 3.
+            norm (str, optional): Normalization type. Defaults to 'batch'.
+            activation (str, optional): Activation function type. Defaults to 'leaky_relu'.
+            leaky_relu (float, optional): Negative slope for LeakyReLU. Defaults to 0.2.
+        '''
+        super().__init__()
+
+        sequence = [
+            ConvBlock(
+                in_channels=in_channels,
+                out_channels=hidden_dim,
+                kernel_size=4,
+                stride=2,
+                padding=1,
+                dim=2,
+                norm=None,
+                activation=activation,
+                leaky_relu=leaky_relu,
+                bias=True
+            )
+        ]
+
+        channel_mult = 1
+        channel_mult_prev = 1
+        for n in range(1, num_layers):
+            channel_mult_prev = channel_mult
+            channel_mult = min(2 ** n, 8)
+            sequence += [
+                ConvBlock(
+                    in_channels=hidden_dim * channel_mult_prev,
+                    out_channels=hidden_dim * channel_mult,
+                    kernel_size=4,
+                    stride=2,
+                    padding=1,
+                    dim=2,
+                    norm=norm,
+                    activation=activation,
+                    leaky_relu=leaky_relu,
+                    bias=False
+                )
+            ]
+
+        channel_mult_prev = channel_mult
+        channel_mult = min(2 ** num_layers, 8)
+
+        sequence += [
+            ConvBlock(
+                in_channels=hidden_dim * channel_mult_prev,
+                out_channels=hidden_dim * channel_mult,
+                kernel_size=4,
+                stride=1,
+                padding=1,
+                dim=2,
+                norm=norm,
+                activation=activation,
+                leaky_relu=leaky_relu,
+                bias=False
+            ),
+            ConvBlock(
+                in_channels=hidden_dim * channel_mult,
+                out_channels=1,
+                kernel_size=4,
+                stride=1,
+                padding=1,
+                dim=2,
+                norm=None,
+                activation=None,
+                bias=True
+            )
+        ]
+
+        self.main = nn.Sequential(*sequence)
+
+    @staticmethod
+    def auto_build(
+        in_channels: int,
+        hidden_dim: int,
+        image_size: int,
+        norm: str = 'batch',
+        activation: str = 'leaky_relu',
+        leaky_relu: float = 0.2
+    ) -> 'PatchDiscriminator':
+        '''
+        Automatically builds a PatchDiscriminator based on the image size.
+
+        Args:
+            in_channels (int): Number of input channels.
+            hidden_dim (int): Base number of filters (channels).
+            image_size (int): Size of the input image.
+            norm (str, optional): Normalization type. Defaults to 'batch'.
+            activation (str, optional): Activation function type. Defaults to 'leaky_relu'.
+            leaky_relu (float, optional): Negative slope for LeakyReLU. Defaults to 0.2.
+
+        Returns:
+            PatchDiscriminator: The constructed PatchDiscriminator.
+        '''
+        num_layers = int(math.log2(image_size / 32))
+        num_layers = max(1, num_layers)
+
+        return PatchDiscriminator(
+            in_channels=in_channels,
+            hidden_dim=hidden_dim,
+            num_layers=num_layers,
+            norm=norm,
+            activation=activation,
+            leaky_relu=leaky_relu
+        )
+
+    def forward(self, input_tensor: torch.Tensor) -> torch.Tensor:
+        '''
+        Defines the computation performed at every call.
+
+        Args:
+            input_tensor (torch.Tensor): The input data.
+
+        Returns:
+            torch.Tensor: The output of the discriminator.
+        '''
+        return self.main(input_tensor)

codon/model/tcn.py

@@ -0,0 +1,124 @@
+from codon.base import *
+from typing     import List, Optional
+
+from codon.block.conv import CausalConv1d, calculate_causal_layer
+
+class TemporalConvNet(BasicModel):
+    '''
+    Temporal Convolutional Network.
+    
+    Consists of a series of Causal Dilated Convolution layers.
+    Supports manually specifying the number of channels for each layer.
+    Use `TemporalConvNet.auto_build` for automatic construction based on the target receptive field.
+    '''
+
+    def __init__(
+        self,
+        in_channels: int,
+        num_channels: List[int],
+        kernel_size: int = 3,
+        dropout: float = 0.2,
+        use_res: bool = True,
+        norm: str = None,
+        activation: str = 'leaky_relu',
+        leaky_relu: float = 0.1,
+        channel_first: bool = True
+    ):
+        '''
+        Initializes the TCN module manually.
+
+        Args:
+            in_channels (int): Number of input channels.
+            num_channels (List[int]): List of output channels for each layer.
+            kernel_size (int, optional): Kernel size. Defaults to 3.
+            dropout (float, optional): Dropout probability. Defaults to 0.2.
+            use_res (bool, optional): Whether to use residual connections. Defaults to True.
+            norm (str, optional): Normalization type (passed to CausalConv1d/ConvBlock). Defaults to None.
+            activation (str, optional): Activation function type. Defaults to 'leaky_relu'.
+            leaky_relu (float, optional): Negative slope for LeakyReLU. Defaults to 0.1.
+            channel_first (bool, optional): Whether input is (Batch, Channels, Seq_Len). Defaults to True.
+        '''
+        super().__init__()
+
+        self.in_channels = in_channels
+        self.kernel_size = kernel_size
+        self.dropout = dropout
+        self.channel_first = channel_first
+        
+        self.network = CausalConv1d.manual_block(
+            in_channels=in_channels,
+            num_channels=num_channels,
+            kernel_size=kernel_size,
+            norm=norm,
+            activation=activation,
+            leaky_relu=leaky_relu,
+            use_res=use_res,
+            dropout=dropout
+        )
+        self.out_channels = num_channels[-1]
+
+    @staticmethod
+    def auto_build(
+        in_channels: int,
+        out_channels: int,
+        receptive_field: int,
+        kernel_size: int = 3,
+        dropout: float = 0.2,
+        use_res: bool = True,
+        norm: str = None,
+        activation: str = 'leaky_relu',
+        leaky_relu: float = 0.1,
+        channel_first: bool = True
+    ) -> 'TemporalConvNet':
+        '''
+        Automatically builds a TCN module based on the target receptive field.
+
+        Args:
+            in_channels (int): Number of input channels.
+            out_channels (int): Unified output channels for each layer.
+            receptive_field (int): Target receptive field (time steps).
+            kernel_size (int, optional): Kernel size. Defaults to 3.
+            dropout (float, optional): Dropout probability. Defaults to 0.2.
+            use_res (bool, optional): Whether to use residual connections. Defaults to True.
+            norm (str, optional): Normalization type. Defaults to None.
+            activation (str, optional): Activation function type. Defaults to 'leaky_relu'.
+            leaky_relu (float, optional): Negative slope for LeakyReLU. Defaults to 0.1.
+            channel_first (bool, optional): Whether input is (Batch, Channels, Seq_Len). Defaults to True.
+
+        Returns:
+            TemporalConvNet: An initialized TCN module.
+        '''
+        layers, _ = calculate_causal_layer(receptive_field, kernel_size)
+        num_channels = [out_channels] * layers
+        
+        return TemporalConvNet(
+            in_channels=in_channels,
+            num_channels=num_channels,
+            kernel_size=kernel_size,
+            dropout=dropout,
+            use_res=use_res,
+            norm=norm,
+            activation=activation,
+            leaky_relu=leaky_relu,
+            channel_first=channel_first
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        '''
+        Forward pass.
+
+        Args:
+            x (torch.Tensor): Input tensor. Shape: [Batch, in_channels, Seq_Len] or [Batch, Seq_Len, in_channels] if channel_first=False.
+
+        Returns:
+            torch.Tensor: Output tensor.
+        '''
+        if not self.channel_first:
+            x = x.transpose(1, 2)
+            
+        x = self.network(x)
+        
+        if not self.channel_first:
+            x = x.transpose(1, 2)
+            
+        return x

codon/ops/__init__.py

@@ -0,0 +1,3 @@
+
+
+__seed__: int = None

codon/ops/attention.py

@@ -0,0 +1,107 @@
+import torch.nn.functional as F
+import torch
+import math
+
+from dataclasses import dataclass
+from typing      import Optional, Tuple
+
+from torch.nn.attention import SDPBackend, sdpa_kernel
+
+@dataclass
+class AttentionOutput:
+    '''
+    Output of the attention mechanism.
+
+    Args:
+        output (torch.Tensor): The output tensor from the attention mechanism.
+        attention_weights (Optional[torch.Tensor], optional): Attention weights. Defaults to None.
+        past_key_value (Optional[Tuple[torch.Tensor, torch.Tensor]], optional): 
+            Cached key and value tensors for autoregressive generation. Defaults to None.
+    '''
+    output: torch.Tensor
+    attention_weights: Optional[torch.Tensor] = None
+    past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None
+
+
+def apply_attention(
+    query_states: torch.Tensor, 
+    key_states: torch.Tensor, 
+    value_states: torch.Tensor, 
+    attention_mask: torch.Tensor = None, 
+    output_attentions: bool = False,
+    is_causal: bool = None,
+    dropout: float = 0.0
+) -> AttentionOutput:
+    ''' 
+    Compute scaled dot-product attention.
+
+    Args:
+        query_states (torch.Tensor): Query states tensor.
+        key_states (torch.Tensor): Key states tensor.
+        value_states (torch.Tensor): Value states tensor.
+        attention_mask (torch.Tensor, optional): Attention mask. Defaults to None.
+        output_attentions (bool, optional): Whether to output attention weights. Defaults to False.
+        is_causal (bool, optional): Whether to apply a causal mask. Defaults to None.
+        dropout (float, optional): Dropout probability. Defaults to 0.0.
+
+    Returns:
+        AttentionOutput: Object containing attention output and optional weights.
+    '''
+    
+    if attention_mask is not None:
+        if attention_mask.dtype != torch.float32:
+            attention_mask = attention_mask.float()
+        
+        if attention_mask.max() <= 1.0:
+            attention_mask = torch.where(attention_mask == 0, float('-inf'), 0.0)
+        
+    if is_causal:
+        tgt_len = query_states.size(-2)
+        src_len = key_states.size(-2)
+        
+        causal_mask = torch.tril(
+            torch.ones((tgt_len, src_len), device=query_states.device, dtype=query_states.dtype)
+        ).view(1, 1, tgt_len, src_len)
+        
+        causal_mask = torch.where(causal_mask == 0, float('-inf'), 0.0)
+        
+        if attention_mask is not None:
+            attention_mask = attention_mask + causal_mask
+        else:
+            attention_mask = causal_mask
+        
+        is_causal = False
+        
+    if not output_attentions:
+        if attention_mask is None and is_causal is None: 
+            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
+                )
+            return AttentionOutput(output=output, attention_weights=None)
+        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)
+
+    if attention_mask is not None:
+        scores = scores + attention_mask
+    attention_weights = torch.softmax(scores, dim=-1)
+    
+    if dropout > 0.0:
+        attention_weights = F.dropout(attention_weights, p=dropout)
+    output = torch.matmul(attention_weights, value_states)
+    
+    return AttentionOutput(output=output, attention_weights=attention_weights)

codon/utils/dataset/__init__.py

@@ -0,0 +1,3 @@
+from .flatdata import FlatDataset, FlatColumnDataset, MappedFlatDataset
+from .corpus   import FileType, CorpusData, CorpusDataset
+from .dataviewer import DataViewer, preview_fields

codon/utils/dataset/base.py

@@ -0,0 +1,46 @@
+from typing import Any
+
+class CodonDataset:
+    '''
+    Base class for all Codon datasets.
+
+    This abstract class defines the interface that all datasets must implement.
+    It provides a common structure for accessing data rows and length.
+    '''
+
+    @property
+    def row(self) -> int:
+        '''
+        Returns the number of rows in the dataset.
+
+        Returns:
+            int: The total number of rows.
+        '''
+        return len(self)
+
+    def __len__(self) -> int:
+        '''
+        Returns the length of the dataset.
+
+        Returns:
+            int: The length of the dataset.
+
+        Raises:
+            NotImplementedError: If the method is not implemented by the subclass.
+        '''
+        raise NotImplementedError
+
+    def __getitem__(self, idx: Any) -> Any:
+        '''
+        Retrieves an item from the dataset at the specified index.
+
+        Args:
+            idx (Any): The index of the item to retrieve.
+
+        Returns:
+            Any: The item at the specified index.
+
+        Raises:
+            NotImplementedError: If the method is not implemented by the subclass.
+        '''
+        raise NotImplementedError