deeplotx 0.5.6__tar.gz → 0.8.0__tar.gz

{deeplotx-0.5.6 → deeplotx-0.8.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: deeplotx
-Version: 0.5.6
+Version: 0.8.0
 Summary: Easy-2-use long text NLP toolkit.
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
@@ -13,7 +13,7 @@ Requires-Dist: python-dotenv
 Requires-Dist: torch
 Requires-Dist: transformers
 Requires-Dist: typing-extensions
-Requires-Dist: vortezwohl>=0.0.6
+Requires-Dist: vortezwohl>=0.0.8
 Dynamic: license-file
 
 [![Ask DeepWiki](https://deepwiki.com/badge.svg)](https://deepwiki.com/vortezwohl/DeepLoTX)
@@ -163,6 +163,8 @@ Dynamic: license-file
 
     ```python
     from deeplotx import (
+        BaseNeuralNetwork,  # 深度神经网络基类
+        FeedForward,  # 前馈神经网络
         LinearRegression,  # 线性回归
         LogisticRegression,  # 逻辑回归 / 二分类 / 多标签分类
         SoftmaxRegression,  # Softmax 回归 / 多分类
@@ -181,38 +183,54 @@ Dynamic: license-file
 
     import torch
     from torch import nn
-    
+
     from deeplotx.nn.base_neural_network import BaseNeuralNetwork
-    
-    
-    class LinearRegression(BaseNeuralNetwork):
-        def __init__(self, input_dim: int, output_dim: int, model_name: str | None = None,
-                     device: str | None = None, dtype: torch.dtype | None = None):
-            super().__init__(model_name=model_name, device=device, dtype=dtype)
-            self.fc1 = nn.Linear(input_dim, 1024, device=self.device, dtype=self.dtype)
-            self.fc1_to_fc4_res = nn.Linear(1024, 64, device=self.device, dtype=self.dtype)
-            self.fc2 = nn.Linear(1024, 768, device=self.device, dtype=self.dtype)
-            self.fc3 = nn.Linear(768, 128, device=self.device, dtype=self.dtype)
-            self.fc4 = nn.Linear(128, 64, device=self.device, dtype=self.dtype)
-            self.fc5 = nn.Linear(64, output_dim, device=self.device, dtype=self.dtype)
-            self.parametric_relu_1 = nn.PReLU(num_parameters=1, init=5e-3, device=self.device, dtype=self.dtype)
-            self.parametric_relu_2 = nn.PReLU(num_parameters=1, init=5e-3, device=self.device, dtype=self.dtype)
-            self.parametric_relu_3 = nn.PReLU(num_parameters=1, init=5e-3, device=self.device, dtype=self.dtype)
-            self.parametric_relu_4 = nn.PReLU(num_parameters=1, init=5e-3, device=self.device, dtype=self.dtype)
-    
+
+
+    class FeedForwardUnit(BaseNeuralNetwork):
+        def __init__(self, feature_dim: int, expansion_factor: int | float = 2,
+                    bias: bool = True, dropout_rate: float = 0.05, model_name: str | None = None,
+                    device: str | None = None, dtype: torch.dtype | None = None):
+            super().__init__(in_features=feature_dim, out_features=feature_dim, model_name=model_name, device=device, dtype=dtype)
+            self._dropout_rate = dropout_rate
+            self.fc1 = nn.Linear(feature_dim, int(feature_dim * expansion_factor), bias=bias,
+                                device=self.device, dtype=self.dtype)
+            self.fc2 = nn.Linear(int(feature_dim * expansion_factor), feature_dim, bias=bias,
+                                device=self.device, dtype=self.dtype)
+            self.parametric_relu_1 = nn.PReLU(num_parameters=1, init=5e-3,
+                                            device=self.device, dtype=self.dtype)
+            self.layer_norm = nn.LayerNorm(normalized_shape=self.fc1.in_features, eps=1e-9,
+                                        device=self.device, dtype=self.dtype)
+
         @override
-        def forward(self, x) -> torch.Tensor:
+        def forward(self, x: torch.Tensor) -> torch.Tensor:
             x = self.ensure_device_and_dtype(x, device=self.device, dtype=self.dtype)
-            fc1_out = self.parametric_relu_1(self.fc1(x))
-            x = nn.LayerNorm(normalized_shape=1024, eps=1e-9, device=self.device, dtype=self.dtype)(fc1_out)
-            x = torch.dropout(x, p=0.2, train=self.training)
-            x = self.parametric_relu_2(self.fc2(x))
-            x = nn.LayerNorm(normalized_shape=768, eps=1e-9, device=self.device, dtype=self.dtype)(x)
-            x = torch.dropout(x, p=0.2, train=self.training)
-            x = self.parametric_relu_3(self.fc3(x))
-            x = torch.dropout(x, p=0.2, train=self.training)
-            x = self.parametric_relu_4(self.fc4(x)) + self.fc1_to_fc4_res(fc1_out)
-            x = self.fc5(x)
+            residual = x
+            x = self.layer_norm(x)
+            x = self.fc1(x)
+            x = self.parametric_relu_1(x)
+            if self._dropout_rate > .0:
+                x = torch.dropout(x, p=self._dropout_rate, train=self.training)
+            return self.fc2(x) + residual
+
+
+    class FeedForward(BaseNeuralNetwork):
+        def __init__(self, feature_dim: int, num_layers: int = 1, expansion_factor: int | float = 2,
+                    bias: bool = True, dropout_rate: float = 0.05, model_name: str | None = None,
+                    device: str | None = None, dtype: torch.dtype | None = None):
+            if num_layers < 1:
+                raise ValueError('num_layers cannot be less than 1.')
+            super().__init__(in_features=feature_dim, out_features=feature_dim, model_name=model_name, device=device, dtype=dtype)
+            self.ffn_layers = nn.ModuleList([FeedForwardUnit(feature_dim=feature_dim,
+                                                            expansion_factor=expansion_factor, bias=bias,
+                                                            dropout_rate=dropout_rate,
+                                                            device=self.device, dtype=self.dtype)] * num_layers)
+
+        @override
+        def forward(self, x: torch.Tensor) -> torch.Tensor:
+            x = self.ensure_device_and_dtype(x, device=self.device, dtype=self.dtype)
+            for ffn in self.ffn_layers:
+                x = ffn(x)
             return x
     ```
 
@@ -222,29 +240,34 @@ Dynamic: license-file
     from typing_extensions import override
 
     import torch
-    from torch import nn, softmax
 
     from deeplotx.nn.base_neural_network import BaseNeuralNetwork
+    from deeplotx.nn.feed_forward import FeedForward
 
 
     class SelfAttention(BaseNeuralNetwork):
-        def __init__(self, feature_dim: int, model_name: str | None = None,
-                    device: str | None = None, dtype: torch.dtype | None = None):
-            super().__init__(model_name=model_name, device=device, dtype=dtype)
+        def __init__(self, feature_dim: int, bias: bool = True, proj_layers: int = 1,
+                    proj_expansion_factor: int | float = 1.5, dropout_rate: float = 0.02,
+                    model_name: str | None = None, device: str | None = None, dtype: torch.dtype | None = None):
+            super().__init__(in_features=feature_dim, out_features=feature_dim, model_name=model_name,
+                            device=device, dtype=dtype)
             self._feature_dim = feature_dim
-            self.q_proj = nn.Linear(in_features=self._feature_dim, out_features=self._feature_dim,
-                                    bias=True, device=self.device, dtype=self.dtype)
-            self.k_proj = nn.Linear(in_features=self._feature_dim, out_features=self._feature_dim,
-                                    bias=True, device=self.device, dtype=self.dtype)
-            self.v_proj = nn.Linear(in_features=self._feature_dim, out_features=self._feature_dim,
-                                    bias=True, device=self.device, dtype=self.dtype)
+            self.q_proj = FeedForward(feature_dim=self._feature_dim, num_layers=proj_layers,
+                                    expansion_factor=proj_expansion_factor,
+                                    bias=bias, dropout_rate=dropout_rate, device=self.device, dtype=self.dtype)
+            self.k_proj = FeedForward(feature_dim=self._feature_dim, num_layers=proj_layers,
+                                    expansion_factor=proj_expansion_factor,
+                                    bias=bias, dropout_rate=dropout_rate, device=self.device, dtype=self.dtype)
+            self.v_proj = FeedForward(feature_dim=self._feature_dim, num_layers=proj_layers,
+                                    expansion_factor=proj_expansion_factor,
+                                    bias=bias, dropout_rate=dropout_rate, device=self.device, dtype=self.dtype)
 
         def _attention(self, x: torch.Tensor, mask: torch.Tensor | None = None) -> torch.Tensor:
             q, k = self.q_proj(x), self.k_proj(x)
             attn = torch.matmul(q, k.transpose(-2, -1))
             attn = attn / (self._feature_dim ** 0.5)
             attn = attn.masked_fill(mask == 0, -1e9) if mask is not None else attn
-            return softmax(attn, dim=-1)
+            return torch.softmax(attn, dim=-1)
 
         @override
         def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None) -> torch.Tensor:

{deeplotx-0.5.6 → deeplotx-0.8.0}/README.md

@@ -145,6 +145,8 @@
 
     ```python
     from deeplotx import (
+        BaseNeuralNetwork,  # 深度神经网络基类
+        FeedForward,  # 前馈神经网络
         LinearRegression,  # 线性回归
         LogisticRegression,  # 逻辑回归 / 二分类 / 多标签分类
         SoftmaxRegression,  # Softmax 回归 / 多分类
@@ -163,38 +165,54 @@
 
     import torch
     from torch import nn
-    
+
     from deeplotx.nn.base_neural_network import BaseNeuralNetwork
-    
-    
-    class LinearRegression(BaseNeuralNetwork):
-        def __init__(self, input_dim: int, output_dim: int, model_name: str | None = None,
-                     device: str | None = None, dtype: torch.dtype | None = None):
-            super().__init__(model_name=model_name, device=device, dtype=dtype)
-            self.fc1 = nn.Linear(input_dim, 1024, device=self.device, dtype=self.dtype)
-            self.fc1_to_fc4_res = nn.Linear(1024, 64, device=self.device, dtype=self.dtype)
-            self.fc2 = nn.Linear(1024, 768, device=self.device, dtype=self.dtype)
-            self.fc3 = nn.Linear(768, 128, device=self.device, dtype=self.dtype)
-            self.fc4 = nn.Linear(128, 64, device=self.device, dtype=self.dtype)
-            self.fc5 = nn.Linear(64, output_dim, device=self.device, dtype=self.dtype)
-            self.parametric_relu_1 = nn.PReLU(num_parameters=1, init=5e-3, device=self.device, dtype=self.dtype)
-            self.parametric_relu_2 = nn.PReLU(num_parameters=1, init=5e-3, device=self.device, dtype=self.dtype)
-            self.parametric_relu_3 = nn.PReLU(num_parameters=1, init=5e-3, device=self.device, dtype=self.dtype)
-            self.parametric_relu_4 = nn.PReLU(num_parameters=1, init=5e-3, device=self.device, dtype=self.dtype)
-    
+
+
+    class FeedForwardUnit(BaseNeuralNetwork):
+        def __init__(self, feature_dim: int, expansion_factor: int | float = 2,
+                    bias: bool = True, dropout_rate: float = 0.05, model_name: str | None = None,
+                    device: str | None = None, dtype: torch.dtype | None = None):
+            super().__init__(in_features=feature_dim, out_features=feature_dim, model_name=model_name, device=device, dtype=dtype)
+            self._dropout_rate = dropout_rate
+            self.fc1 = nn.Linear(feature_dim, int(feature_dim * expansion_factor), bias=bias,
+                                device=self.device, dtype=self.dtype)
+            self.fc2 = nn.Linear(int(feature_dim * expansion_factor), feature_dim, bias=bias,
+                                device=self.device, dtype=self.dtype)
+            self.parametric_relu_1 = nn.PReLU(num_parameters=1, init=5e-3,
+                                            device=self.device, dtype=self.dtype)
+            self.layer_norm = nn.LayerNorm(normalized_shape=self.fc1.in_features, eps=1e-9,
+                                        device=self.device, dtype=self.dtype)
+
         @override
-        def forward(self, x) -> torch.Tensor:
+        def forward(self, x: torch.Tensor) -> torch.Tensor:
             x = self.ensure_device_and_dtype(x, device=self.device, dtype=self.dtype)
-            fc1_out = self.parametric_relu_1(self.fc1(x))
-            x = nn.LayerNorm(normalized_shape=1024, eps=1e-9, device=self.device, dtype=self.dtype)(fc1_out)
-            x = torch.dropout(x, p=0.2, train=self.training)
-            x = self.parametric_relu_2(self.fc2(x))
-            x = nn.LayerNorm(normalized_shape=768, eps=1e-9, device=self.device, dtype=self.dtype)(x)
-            x = torch.dropout(x, p=0.2, train=self.training)
-            x = self.parametric_relu_3(self.fc3(x))
-            x = torch.dropout(x, p=0.2, train=self.training)
-            x = self.parametric_relu_4(self.fc4(x)) + self.fc1_to_fc4_res(fc1_out)
-            x = self.fc5(x)
+            residual = x
+            x = self.layer_norm(x)
+            x = self.fc1(x)
+            x = self.parametric_relu_1(x)
+            if self._dropout_rate > .0:
+                x = torch.dropout(x, p=self._dropout_rate, train=self.training)
+            return self.fc2(x) + residual
+
+
+    class FeedForward(BaseNeuralNetwork):
+        def __init__(self, feature_dim: int, num_layers: int = 1, expansion_factor: int | float = 2,
+                    bias: bool = True, dropout_rate: float = 0.05, model_name: str | None = None,
+                    device: str | None = None, dtype: torch.dtype | None = None):
+            if num_layers < 1:
+                raise ValueError('num_layers cannot be less than 1.')
+            super().__init__(in_features=feature_dim, out_features=feature_dim, model_name=model_name, device=device, dtype=dtype)
+            self.ffn_layers = nn.ModuleList([FeedForwardUnit(feature_dim=feature_dim,
+                                                            expansion_factor=expansion_factor, bias=bias,
+                                                            dropout_rate=dropout_rate,
+                                                            device=self.device, dtype=self.dtype)] * num_layers)
+
+        @override
+        def forward(self, x: torch.Tensor) -> torch.Tensor:
+            x = self.ensure_device_and_dtype(x, device=self.device, dtype=self.dtype)
+            for ffn in self.ffn_layers:
+                x = ffn(x)
             return x
     ```
 
@@ -204,29 +222,34 @@
     from typing_extensions import override
 
     import torch
-    from torch import nn, softmax
 
     from deeplotx.nn.base_neural_network import BaseNeuralNetwork
+    from deeplotx.nn.feed_forward import FeedForward
 
 
     class SelfAttention(BaseNeuralNetwork):
-        def __init__(self, feature_dim: int, model_name: str | None = None,
-                    device: str | None = None, dtype: torch.dtype | None = None):
-            super().__init__(model_name=model_name, device=device, dtype=dtype)
+        def __init__(self, feature_dim: int, bias: bool = True, proj_layers: int = 1,
+                    proj_expansion_factor: int | float = 1.5, dropout_rate: float = 0.02,
+                    model_name: str | None = None, device: str | None = None, dtype: torch.dtype | None = None):
+            super().__init__(in_features=feature_dim, out_features=feature_dim, model_name=model_name,
+                            device=device, dtype=dtype)
             self._feature_dim = feature_dim
-            self.q_proj = nn.Linear(in_features=self._feature_dim, out_features=self._feature_dim,
-                                    bias=True, device=self.device, dtype=self.dtype)
-            self.k_proj = nn.Linear(in_features=self._feature_dim, out_features=self._feature_dim,
-                                    bias=True, device=self.device, dtype=self.dtype)
-            self.v_proj = nn.Linear(in_features=self._feature_dim, out_features=self._feature_dim,
-                                    bias=True, device=self.device, dtype=self.dtype)
+            self.q_proj = FeedForward(feature_dim=self._feature_dim, num_layers=proj_layers,
+                                    expansion_factor=proj_expansion_factor,
+                                    bias=bias, dropout_rate=dropout_rate, device=self.device, dtype=self.dtype)
+            self.k_proj = FeedForward(feature_dim=self._feature_dim, num_layers=proj_layers,
+                                    expansion_factor=proj_expansion_factor,
+                                    bias=bias, dropout_rate=dropout_rate, device=self.device, dtype=self.dtype)
+            self.v_proj = FeedForward(feature_dim=self._feature_dim, num_layers=proj_layers,
+                                    expansion_factor=proj_expansion_factor,
+                                    bias=bias, dropout_rate=dropout_rate, device=self.device, dtype=self.dtype)
 
         def _attention(self, x: torch.Tensor, mask: torch.Tensor | None = None) -> torch.Tensor:
             q, k = self.q_proj(x), self.k_proj(x)
             attn = torch.matmul(q, k.transpose(-2, -1))
             attn = attn / (self._feature_dim ** 0.5)
             attn = attn.masked_fill(mask == 0, -1e9) if mask is not None else attn
-            return softmax(attn, dim=-1)
+            return torch.softmax(attn, dim=-1)
 
         @override
         def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None) -> torch.Tensor:

{deeplotx-0.5.6 → deeplotx-0.8.0}/deeplotx/__init__.py

@@ -5,12 +5,16 @@ __ROOT__ = os.path.dirname(os.path.abspath(__file__))
 
 from .encoder import Encoder, LongTextEncoder, LongformerEncoder
 from .nn import (
+    FeedForward,
     LinearRegression,
     LogisticRegression,
     SoftmaxRegression,
     RecursiveSequential,
     LongContextRecursiveSequential,
-    SelfAttention,
+    RoPE,
+    Attention,
+    MultiHeadAttention,
+    RoFormerEncoder,
     AutoRegression,
     LongContextAutoRegression
 )

deeplotx-0.8.0/deeplotx/encoder/encoder.py

@@ -0,0 +1,66 @@
+import logging
+import os
+import math
+from requests.exceptions import ConnectTimeout, SSLError
+
+import torch
+from torch import nn
+from transformers import AutoTokenizer, AutoModel
+
+from deeplotx import __ROOT__
+
+CACHE_PATH = os.path.join(__ROOT__, '.cache')
+DEFAULT_BERT = 'FacebookAI/xlm-roberta-base'
+logger = logging.getLogger('deeplotx.embedding')
+
+
+class Encoder(nn.Module):
+    def __init__(self, model_name_or_path: str = DEFAULT_BERT, device: str | None = None):
+        super().__init__()
+        self.device = torch.device(device) if device is not None \
+            else torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        try:
+            self.tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path=model_name_or_path,
+                                                           cache_dir=CACHE_PATH, _from_auto=True,
+                                                           trust_remote_code=True)
+            self.encoder = AutoModel.from_pretrained(pretrained_model_name_or_path=model_name_or_path,
+                                                     cache_dir=CACHE_PATH, _from_auto=True,
+                                                     trust_remote_code=True).to(self.device)
+        except ConnectTimeout:
+            self.tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path=model_name_or_path,
+                                                           cache_dir=CACHE_PATH, _from_auto=True,
+                                                           trust_remote_code=True, local_files_only=True)
+            self.encoder = AutoModel.from_pretrained(pretrained_model_name_or_path=model_name_or_path,
+                                                     cache_dir=CACHE_PATH, _from_auto=True,
+                                                     trust_remote_code=True, local_files_only=True).to(self.device)
+        except SSLError:
+            self.tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path=model_name_or_path,
+                                                           cache_dir=CACHE_PATH, _from_auto=True,
+                                                           trust_remote_code=True, local_files_only=True)
+            self.encoder = AutoModel.from_pretrained(pretrained_model_name_or_path=model_name_or_path,
+                                                     cache_dir=CACHE_PATH, _from_auto=True,
+                                                     trust_remote_code=True, local_files_only=True).to(self.device)
+        self.embed_dim = self.encoder.config.max_position_embeddings
+        logger.debug(f'{Encoder.__name__} initialized on device: {self.device}.')
+
+    def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor, *args, **kwargs) -> torch.Tensor:
+        def _encoder(_input_tup: tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor:
+            return self.encoder.forward(_input_tup[0], attention_mask=_input_tup[1]).last_hidden_state[:, 0, :]
+
+        num_chunks = math.ceil(input_ids.shape[-1] / self.embed_dim)
+        chunks = chunk_results = []
+        for i in range(num_chunks):
+            start_idx = i * self.embed_dim
+            end_idx = min(start_idx + self.embed_dim, input_ids.shape[-1])
+            chunks.append((input_ids[:, start_idx: end_idx], attention_mask[:, start_idx: end_idx]))
+        ori_mode = self.encoder.training
+        self.encoder.eval()
+        with torch.no_grad():
+            chunk_results = [_encoder(x) for x in chunks]
+        self.encoder.train(mode=ori_mode)
+        return torch.cat(chunk_results, dim=-1)
+
+    def encode(self, text: str) -> torch.Tensor:
+        _input_ids = torch.tensor([self.tokenizer.encode(text)], dtype=torch.long, device=self.device)
+        _att_mask = torch.tensor([[1] * _input_ids.shape[-1]], dtype=torch.int, device=self.device)
+        return self.forward(_input_ids, _att_mask).squeeze()

{deeplotx-0.5.6 → deeplotx-0.8.0}/deeplotx/encoder/long_text_encoder.py

@@ -15,12 +15,14 @@ logger = logging.getLogger('deeplotx.embedding')
 class LongTextEncoder(Encoder):
     def __init__(self, max_length: int, chunk_size: int = 448,
                  overlapping: int = 32, model_name_or_path: str = DEFAULT_BERT,
-                 cache_capacity: int = 64, device: str | None = None):
+                 cache_capacity: int = 64, max_workers: int = 8, device: str | None = None):
         super().__init__(model_name_or_path=model_name_or_path, device=device)
+        assert overlapping < chunk_size, f'overlapping ({overlapping}) must be less than chunk size ({chunk_size}).'
         self._max_length = max_length
         self._chunk_size = chunk_size
         self._overlapping = overlapping
         self._cache = LRUCache(capacity=cache_capacity)
+        self._worker_group = ThreadPool(max_workers=max_workers)
 
     def __chunk_embedding(self, idx: int, x: torch.Tensor, mask: torch.Tensor) -> tuple[int, torch.Tensor]:
         return idx, super().forward(x, attention_mask=mask)
@@ -63,7 +65,7 @@ class LongTextEncoder(Encoder):
             _tmp_right = (i + 1) * self._chunk_size + self._overlapping
             chunks.append((i, torch.tensor([_text_to_input_ids[_tmp_left: _tmp_right]], dtype=torch.int, device=self.device),
                            torch.tensor([_text_to_input_ids_att_mask[_tmp_left: _tmp_right]], dtype=torch.int, device=self.device)))
-        embeddings = list(ThreadPool(max_workers=min(num_chunks + 1, 8)).map(self.__chunk_embedding, chunks))
+        embeddings = list(self._worker_group.map(self.__chunk_embedding, chunks))
         embeddings = sorted([x.returns for x in embeddings], key=lambda x: x[0], reverse=False)
         fin_embedding = [x[1] for x in embeddings]
         # write cache

{deeplotx-0.5.6 → deeplotx-0.8.0}/deeplotx/nn/__init__.py

@@ -1,8 +1,13 @@
+from .base_neural_network import BaseNeuralNetwork
+from .feed_forward import FeedForward
 from .linear_regression import LinearRegression
 from .logistic_regression import LogisticRegression
 from .softmax_regression import SoftmaxRegression
 from .recursive_sequential import RecursiveSequential
 from .long_context_recursive_sequential import LongContextRecursiveSequential
-from .self_attention import SelfAttention
+from .rope import RoPE
+from .attention import Attention
+from .multi_head_attention import MultiHeadAttention
+from .roformer_encoder import RoFormerEncoder
 from .auto_regression import AutoRegression
 from .long_context_auto_regression import LongContextAutoRegression

deeplotx-0.8.0/deeplotx/nn/attention.py

@@ -0,0 +1,48 @@
+from typing_extensions import override
+
+import torch
+
+from deeplotx.nn.base_neural_network import BaseNeuralNetwork
+from deeplotx.nn.feed_forward import FeedForward
+from deeplotx.nn.rope import RoPE, DEFAULT_THETA
+
+
+class Attention(BaseNeuralNetwork):
+    def __init__(self, feature_dim: int, bias: bool = True, positional: bool = True,
+                 proj_layers: int = 1, proj_expansion_factor: int | float = 1.5, dropout_rate: float = 0.02,
+                 model_name: str | None = None, device: str | None = None, dtype: torch.dtype | None = None,
+                 **kwargs):
+        super().__init__(in_features=feature_dim, out_features=feature_dim, model_name=model_name,
+                         device=device, dtype=dtype)
+        self._positional = positional
+        self._feature_dim = feature_dim
+        self.q_proj = FeedForward(feature_dim=self._feature_dim, num_layers=proj_layers,
+                                  expansion_factor=proj_expansion_factor,
+                                  bias=bias, dropout_rate=dropout_rate, device=self.device, dtype=self.dtype)
+        self.k_proj = FeedForward(feature_dim=self._feature_dim, num_layers=proj_layers,
+                                  expansion_factor=proj_expansion_factor,
+                                  bias=bias, dropout_rate=dropout_rate, device=self.device, dtype=self.dtype)
+        self.v_proj = FeedForward(feature_dim=self._feature_dim, num_layers=proj_layers,
+                                  expansion_factor=proj_expansion_factor,
+                                  bias=bias, dropout_rate=dropout_rate, device=self.device, dtype=self.dtype)
+        if self._positional:
+            self.rope = RoPE(feature_dim=self._feature_dim, theta=kwargs.get('theta', DEFAULT_THETA),
+                             device=self.device, dtype=self.dtype)
+
+    def _attention(self, x: torch.Tensor, y: torch.Tensor, mask: torch.Tensor | None = None) -> torch.Tensor:
+        q, k = self.q_proj(x), self.k_proj(y)
+        if self._positional:
+            q, k = self.rope(q), self.rope(k)
+        attn = torch.matmul(q, k.transpose(-2, -1))
+        attn = attn / (self._feature_dim ** 0.5)
+        attn = attn.masked_fill(mask == 0, -1e9) if mask is not None else attn
+        return torch.softmax(attn, dim=-1)
+
+    @override
+    def forward(self, x: torch.Tensor, y: torch.Tensor | None = None, mask: torch.Tensor | None = None) -> torch.Tensor:
+        x = self.ensure_device_and_dtype(x, device=self.device, dtype=self.dtype)
+        y = x if y is None else self.ensure_device_and_dtype(y, device=self.device, dtype=self.dtype)
+        if mask is not None:
+            mask = self.ensure_device_and_dtype(mask, device=self.device, dtype=self.dtype)
+        v = self.v_proj(y)
+        return torch.matmul(self._attention(x, y, mask), v)

deeplotx-0.8.0/deeplotx/nn/auto_regression.py

@@ -0,0 +1,14 @@
+import torch
+
+from deeplotx.nn import RecursiveSequential
+
+
+class AutoRegression(RecursiveSequential):
+    def __init__(self, feature_dim: int, bias: bool = True,
+                 recursive_layers: int = 1, recursive_hidden_dim: int | None = None,
+                 ffn_layers: int = 1, ffn_expansion_factor: int | float = 2, dropout_rate: float = 0.05,
+                 model_name: str | None = None, device: str | None = None, dtype: torch.dtype | None = None):
+        super().__init__(input_dim=feature_dim, output_dim=feature_dim, bias=bias,
+                         recursive_layers=recursive_layers, recursive_hidden_dim=recursive_hidden_dim,
+                         ffn_layers=ffn_layers, ffn_expansion_factor=ffn_expansion_factor,
+                         dropout_rate=dropout_rate, model_name=model_name, device=device, dtype=dtype)

deeplotx-0.8.0/deeplotx/nn/base_neural_network.py

@@ -0,0 +1,140 @@
+import os
+from abc import abstractmethod
+
+import torch
+from torch import nn
+from torch.nn import init
+
+DEFAULT_SUFFIX = 'dlx'
+
+
+class BaseNeuralNetwork(nn.Module):
+    def __init__(self, in_features: int, out_features: int, model_name: str | None = None,
+                 device: str | None = None, dtype: torch.dtype | None = None):
+        super().__init__()
+        self._model_name = model_name \
+            if model_name is not None \
+            else self.__class__.__name__
+        self.device = torch.device(device) if device is not None \
+            else torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        self.dtype = dtype if dtype is not None else torch.float32
+        self._in_features = in_features
+        self._out_features = out_features
+
+    @property
+    def in_features(self) -> int:
+        return self._in_features
+
+    @property
+    def out_features(self) -> int:
+        return self._out_features
+
+    @staticmethod
+    def ensure_device_and_dtype(x: torch.Tensor, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
+        if x.device != device:
+            x = x.to(device)
+        if x.dtype != dtype:
+            x = x.to(dtype)
+        return x
+
+    def initialize_weights(self):
+        for m in self.modules():
+            match m.__class__:
+                case nn.Linear:
+                    init.kaiming_normal_(m.weight, mode='fan_in', nonlinearity='leaky_relu')
+                    if m.bias is not None:
+                        init.constant_(m.bias, 0)
+                case nn.BatchNorm2d | nn.BatchNorm1d | nn.BatchNorm3d:
+                    init.constant_(m.weight, 1)
+                    init.constant_(m.bias, 0)
+                case nn.LSTM | nn.GRU:
+                    for name, param in m.named_parameters():
+                        _tmp_name = name.lower()
+                        if 'weight_ih' in _tmp_name:
+                            init.kaiming_normal_(param, mode='fan_in', nonlinearity='sigmoid')
+                        elif 'weight_hh' in _tmp_name:
+                            init.orthogonal_(param)
+                        elif 'bias' in _tmp_name:
+                            init.constant_(param, 0)
+                case _:
+                    pass
+        return self
+
+    def size(self) -> dict:
+        total_params = trainable_params = non_trainable_params = 0
+        for param in self.parameters():
+            params = param.numel()
+            total_params += params
+            if param.requires_grad:
+                trainable_params += params
+            else:
+                non_trainable_params += params
+        return {
+            'total': total_params,
+            'trainable': trainable_params,
+            'non_trainable': non_trainable_params
+        }
+
+    def l1(self, _lambda: float = 1e-4) -> torch.Tensor:
+        def _l1() -> torch.Tensor:
+            l2_reg = torch.tensor(0., device=self.device, dtype=self.dtype)
+            for param in self.parameters():
+                l2_reg += (torch.abs(param)).sum()
+            return l2_reg
+        return _lambda * _l1()
+
+    def l2(self, _lambda: float = 1e-4) -> torch.Tensor:
+        def _l2() -> torch.Tensor:
+            l2_reg = torch.tensor(0., device=self.device, dtype=self.dtype)
+            for param in self.parameters():
+                l2_reg += (torch.pow(param, exponent=2.)).sum()
+            return l2_reg
+        return _lambda * _l2() / 2.
+
+    def elastic_net(self, alpha: float = 1e-4, rho: float = 0.5) -> torch.Tensor:
+        return alpha * (rho * self.l1(_lambda=1.) + (1 - rho) * self.l2(_lambda=1.))
+
+    @abstractmethod
+    def forward(self, *args, **kwargs) -> torch.Tensor: ...
+
+    def predict(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.ensure_device_and_dtype(x, device=self.device, dtype=self.dtype)
+        __train = self.training
+        self.training = False
+        with torch.no_grad():
+            res = self.forward(x)
+        self.training = __train
+        return res
+
+    def save(self, model_name: str | None = None, model_dir: str = '.', _suffix: str = DEFAULT_SUFFIX):
+        os.makedirs(model_dir, exist_ok=True)
+        model_file_name = f'{model_name}.{_suffix}' if model_name is not None else f'{self._model_name}.{_suffix}'
+        torch.save(self.state_dict(), os.path.join(model_dir, model_file_name))
+        return self
+
+    def load(self, model_name: str | None = None, model_dir: str = '.', _suffix: str = DEFAULT_SUFFIX):
+        model_file_name = f'{model_name}.{_suffix}' if model_name is not None else f'{self._model_name}.{_suffix}'
+        self.load_state_dict(torch.load(os.path.join(model_dir, model_file_name), map_location=self.device, weights_only=True))
+        return self
+
+    def __str__(self):
+        formatted = super().__str__()
+        _line_len = len([sorted(formatted.splitlines(), key=lambda _: len(_), reverse=True)][0])
+        _splitter_1 = '=' * (_line_len + 10)
+        _splitter_2 = '-' * (_line_len + 10)
+        _size = self.size()
+        total_param = _size['total']
+        trainable_param = _size['trainable']
+        non_trainable_param = _size['non_trainable']
+        formatted = (f'{_splitter_1}\n'
+                     f'Model_Name: {self._model_name}\n'
+                     f'In_Features: {self.in_features}\n'
+                     f'Out_Features: {self.out_features}\n'
+                     f'Device: {self.device}\n'
+                     f'Dtype: {self.dtype}\n'
+                     f'Total_Parameters: {total_param}\n'
+                     f'Trainable_Parameters: {trainable_param}\n'
+                     f'NonTrainable_Parameters: {non_trainable_param}\n'
+                     f'{_splitter_2}'
+                     f'\n{formatted}\n{_splitter_1}')
+        return formatted