deeplotx 0.5.5__py3-none-any.whl → 0.6.1__py3-none-any.whl

deeplotx/__init__.py

@@ -5,6 +5,8 @@ __ROOT__ = os.path.dirname(os.path.abspath(__file__))
 
 from .encoder import Encoder, LongTextEncoder, LongformerEncoder
 from .nn import (
+    BaseNeuralNetwork,
+    FeedForward,
     LinearRegression,
     LogisticRegression,
     SoftmaxRegression,

deeplotx/encoder/encoder.py

@@ -1,6 +1,7 @@
 import logging
 import os
 import math
+from requests.exceptions import ConnectTimeout, SSLError
 
 import torch
 from torch import nn
@@ -18,14 +19,31 @@ class Encoder(nn.Module):
         super().__init__()
         self.device = torch.device(device) if device is not None \
             else torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-        self.tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path=model_name_or_path,
-                                                       cache_dir=CACHE_PATH, _from_auto=True)
-        self.encoder = AutoModel.from_pretrained(pretrained_model_name_or_path=model_name_or_path,
-                                                 cache_dir=CACHE_PATH, _from_auto=True).to(self.device)
+        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) -> torch.Tensor:
+    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, :]
 

deeplotx/encoder/long_text_encoder.py

@@ -15,16 +15,21 @@ 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)
         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)
 
+    @override
+    def forward(self, text: str, flatten: bool = False, *args, **kwargs) -> torch.Tensor:
+        return self.encode(text=text, flatten=flatten)
+
     @override
     def encode(self, text: str, flatten: bool = False) -> torch.Tensor:
         def postprocess(tensors: list[torch.Tensor], _flatten: bool) -> torch.Tensor:
@@ -59,7 +64,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/nn/__init__.py

@@ -1,3 +1,5 @@
+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

deeplotx/nn/auto_regression.py

@@ -5,8 +5,11 @@ from deeplotx.nn import RecursiveSequential
 
 class AutoRegression(RecursiveSequential):
     def __init__(self, feature_dim: int, hidden_dim: int | None = None,
-                 recursive_layers: int = 2, model_name: str | None = None,
+                 recursive_layers: int = 2, ffn_layers: int = 1, ffn_expansion_factor: int | float = 2,
+                 ffn_bias: bool = True, ffn_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,
                          hidden_dim=hidden_dim, recursive_layers=recursive_layers,
+                         ffn_layers=ffn_layers, ffn_expansion_factor=ffn_expansion_factor,
+                         ffn_bias=ffn_bias, ffn_dropout_rate=ffn_dropout_rate,
                          model_name=model_name, device=device, dtype=dtype)

deeplotx/nn/base_neural_network.py

@@ -8,7 +8,8 @@ DEFAULT_SUFFIX = 'dlx'
 
 
 class BaseNeuralNetwork(nn.Module):
-    def __init__(self, model_name: str | None = None, device: str | None = None, dtype: torch.dtype | None = None):
+    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 \
@@ -16,6 +17,16 @@ class BaseNeuralNetwork(nn.Module):
         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:

deeplotx/nn/feed_forward.py

@@ -0,0 +1,53 @@
+from typing_extensions import override
+
+import torch
+from torch import nn
+
+from deeplotx.nn.base_neural_network import BaseNeuralNetwork
+
+
+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) -> torch.Tensor:
+        x = self.ensure_device_and_dtype(x, device=self.device, dtype=self.dtype)
+        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

deeplotx/nn/linear_regression.py

@@ -4,34 +4,22 @@ import torch
 from torch import nn
 
 from deeplotx.nn.base_neural_network import BaseNeuralNetwork
+from deeplotx.nn.feed_forward import FeedForward
 
 
 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)
+    def __init__(self, input_dim: int, output_dim: int, num_layers: int = 1,
+                 expansion_factor: int | float = 1.5, bias: bool = True, dropout_rate: float = 0.1,
+                 model_name: str | None = None, device: str | None = None, dtype: torch.dtype | None = None):
+        super().__init__(in_features=input_dim, out_features=output_dim, model_name=model_name, device=device, dtype=dtype)
+        self.ffn = FeedForward(feature_dim=input_dim, num_layers=num_layers, expansion_factor=expansion_factor,
+                               bias=bias, dropout_rate=dropout_rate, device=self.device, dtype=self.dtype)
+        self.proj = nn.Linear(in_features=input_dim, out_features=output_dim,
+                              bias=bias, device=self.device, dtype=self.dtype)
 
     @override
     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)
-        return x
+        residual = x
+        x = self.ffn(x) + residual
+        return self.proj(x)

deeplotx/nn/logistic_regression.py

@@ -6,9 +6,12 @@ from deeplotx.nn.linear_regression import LinearRegression
 
 
 class LogisticRegression(LinearRegression):
-    def __init__(self, input_dim: int, output_dim: int = 1, model_name: str | None = None,
+    def __init__(self, input_dim: int, output_dim: int = 1, num_layers: int = 1, expansion_factor: int | float = 1.5,
+                 bias: bool = True, dropout_rate: float = 0.1, model_name: str | None = None,
                  device: str | None = None, dtype: torch.dtype | None = None):
-        super().__init__(input_dim=input_dim, output_dim=output_dim, model_name=model_name, device=device, dtype=dtype)
+        super().__init__(input_dim=input_dim, output_dim=output_dim, num_layers=num_layers,
+                         expansion_factor=expansion_factor, bias=bias, dropout_rate=dropout_rate,
+                         model_name=model_name, device=device, dtype=dtype)
 
     @override
     def forward(self, x: torch.Tensor) -> torch.Tensor:

deeplotx/nn/long_context_auto_regression.py

@@ -5,8 +5,11 @@ from deeplotx.nn import LongContextRecursiveSequential
 
 class LongContextAutoRegression(LongContextRecursiveSequential):
     def __init__(self, feature_dim: int, hidden_dim: int | None = None,
-                 recursive_layers: int = 2, model_name: str | None = None,
+                 recursive_layers: int = 2, ffn_layers: int = 1, ffn_expansion_factor: int | float = 2,
+                 ffn_bias: bool = True, ffn_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,
                          hidden_dim=hidden_dim, recursive_layers=recursive_layers,
+                         ffn_layers=ffn_layers, ffn_expansion_factor=ffn_expansion_factor,
+                         ffn_bias=ffn_bias, ffn_dropout_rate=ffn_dropout_rate,
                          model_name=model_name, device=device, dtype=dtype)

deeplotx/nn/long_context_recursive_sequential.py

@@ -10,19 +10,25 @@ from deeplotx.nn.self_attention import SelfAttention
 class LongContextRecursiveSequential(RecursiveSequential):
     def __init__(self, input_dim: int, output_dim: int,
                  hidden_dim: int | None = None, recursive_layers: int = 2,
+                 ffn_layers: int = 1, ffn_expansion_factor: int | float = 2,
+                 ffn_bias: bool = True, ffn_dropout_rate: float = 0.05,
                  model_name: str | None = None, device: str | None = None,
-                 dtype: torch.dtype | None = None):
+                 dtype: torch.dtype | None = None, **kwargs):
         super().__init__(input_dim=input_dim, output_dim=output_dim,
                          hidden_dim=hidden_dim, recursive_layers=recursive_layers,
+                         ffn_layers=ffn_layers, ffn_expansion_factor=ffn_expansion_factor,
+                         ffn_bias=ffn_bias, ffn_dropout_rate=ffn_dropout_rate,
                          model_name=model_name, device=device, dtype=dtype)
-        self._feature_dim = input_dim
-        self.self_attention = SelfAttention(feature_dim=input_dim)
-        self.proj = nn.Linear(in_features=input_dim * 2, out_features=input_dim,
-                              bias=True, device=self.device, dtype=self.dtype)
+        self.self_attention = SelfAttention(feature_dim=input_dim, bias=kwargs.get('attn_proj_bias', ffn_bias),
+                                            proj_layers=kwargs.get('attn_proj_layers', 1),
+                                            proj_expansion_factor=kwargs.get('attn_proj_expansion_factor', ffn_expansion_factor),
+                                            dropout_rate=kwargs.get('attn_proj_dropout_rate', ffn_dropout_rate))
+        self.__proj = nn.Linear(in_features=input_dim * 2, out_features=input_dim,
+                                bias=ffn_bias, device=self.device, dtype=self.dtype)
 
     @override
     def forward(self, x: torch.Tensor, state: tuple[torch.Tensor, torch.Tensor]) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
         x = self.ensure_device_and_dtype(x, device=self.device, dtype=self.dtype)
         x = torch.cat([self.self_attention(x), x], dim=-1)
         x = nn.LayerNorm(normalized_shape=x.shape[-1], eps=1e-9, device=self.device, dtype=self.dtype)(x)
-        return super().forward(self.proj(x), state)
+        return super().forward(self.__proj(x), state)

deeplotx/nn/recursive_sequential.py

@@ -4,23 +4,27 @@ import torch
 from torch import nn
 
 from deeplotx.nn.base_neural_network import BaseNeuralNetwork
-from deeplotx.nn import LinearRegression
+from deeplotx.nn.feed_forward import FeedForward
 
 
 class RecursiveSequential(BaseNeuralNetwork):
     def __init__(self, input_dim: int, output_dim: int,
                  hidden_dim: int | None = None, recursive_layers: int = 2,
+                 ffn_layers: int = 1, ffn_expansion_factor: int | float = 2,
+                 ffn_bias: bool = True, ffn_dropout_rate: float = 0.05,
                  model_name: str | None = None, device: str | None = None,
                  dtype: torch.dtype | None = None):
-        super().__init__(model_name=model_name, device=device, dtype=dtype)
+        super().__init__(in_features=input_dim, out_features=output_dim, model_name=model_name, device=device, dtype=dtype)
         if hidden_dim is None:
             hidden_dim = input_dim
         self.lstm = nn.LSTM(input_size=input_dim, hidden_size=hidden_dim,
                             num_layers=recursive_layers, batch_first=True,
                             bias=True, bidirectional=True, device=self.device,
                             dtype=self.dtype)
-        self.regressive_head = LinearRegression(input_dim=hidden_dim * 2, output_dim=output_dim,
-                                                device=self.device, dtype=self.dtype)
+        self.ffn = FeedForward(feature_dim=hidden_dim * 2, num_layers=ffn_layers, expansion_factor=ffn_expansion_factor,
+                               bias=ffn_bias, dropout_rate=ffn_dropout_rate, device=self.device, dtype=self.dtype)
+        self.__proj = nn.Linear(in_features=hidden_dim * 2, out_features=output_dim, bias=ffn_bias,
+                                device=self.device, dtype=self.dtype)
 
     def initial_state(self, batch_size: int = 1) -> tuple[torch.Tensor, torch.Tensor]:
         zeros = torch.zeros(self.lstm.num_layers * 2, batch_size, self.lstm.hidden_size, device=self.device, dtype=self.dtype)
@@ -32,7 +36,10 @@ class RecursiveSequential(BaseNeuralNetwork):
         state = (self.ensure_device_and_dtype(state[0], device=self.device, dtype=self.dtype),
                  self.ensure_device_and_dtype(state[1], device=self.device, dtype=self.dtype))
         x, (hidden_state, cell_state) = self.lstm(x, state)
-        x = self.regressive_head(x[:, -1, :])
+        x = x[:, -1, :]
+        residual = x
+        x = self.ffn(x) + residual
+        x = self.__proj(x)
         return x, (hidden_state, cell_state)
 
     @override

deeplotx/nn/self_attention.py

@@ -1,29 +1,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/nn/softmax_regression.py

@@ -6,9 +6,12 @@ from deeplotx.nn.linear_regression import LinearRegression
 
 
 class SoftmaxRegression(LinearRegression):
-    def __init__(self, input_dim: int, output_dim: int, model_name: str | None = None,
+    def __init__(self, input_dim: int, output_dim: int, num_layers: int = 1, expansion_factor: int | float = 1.5,
+                 bias: bool = True, dropout_rate: float = 0.1, model_name: str | None = None,
                  device: str | None = None, dtype: torch.dtype | None = None):
-        super().__init__(input_dim=input_dim, output_dim=output_dim, model_name=model_name, device=device, dtype=dtype)
+        super().__init__(input_dim=input_dim, output_dim=output_dim, num_layers=num_layers,
+                         expansion_factor=expansion_factor, bias=bias, dropout_rate=dropout_rate,
+                         model_name=model_name, device=device, dtype=dtype)
 
     @override
     def forward(self, x: torch.Tensor) -> torch.Tensor:

deeplotx/trainer/text_binary_classification_trainer.py

@@ -25,13 +25,13 @@ class TextBinaryClassifierTrainer(BaseTrainer):
               num_epochs: int, learning_rate: float = 2e-6, balancing_dataset: bool = True,
               train_loss_threshold: float = 0.0, valid_loss_threshold: float = 0.0,
               alpha: float = 1e-4, rho: float = 0.2,
-              hidden_dim: int = 256, recursive_layers: int = 2) -> LongContextRecursiveSequential:
+              hidden_dim: int = 256, recursive_layers: int = 2, **kwargs) -> LongContextRecursiveSequential:
         if balancing_dataset:
             min_length = min(len(positive_texts), len(negative_texts))
             positive_texts = positive_texts[:min_length]
             negative_texts = negative_texts[:min_length]
         all_texts = positive_texts + negative_texts
-        text_embeddings = [self._long_text_encoder.encode(x, flatten=False, use_cache=True) for x in all_texts]
+        text_embeddings = [self._long_text_encoder.encode(x, flatten=False) for x in all_texts]
         feature_dim = text_embeddings[0].shape[-1]
         dtype = text_embeddings[0].dtype
         labels = ([torch.tensor([1.], dtype=dtype, device=self.device) for _ in range(len(positive_texts))]
@@ -44,15 +44,27 @@ class TextBinaryClassifierTrainer(BaseTrainer):
         valid_dataset = TensorDataset(inputs[train_size:], labels[train_size:])
         self.train_dataset_loader = DataLoader(train_dataset, batch_size=self._batch_size, shuffle=True)
         self.valid_dataset_loader = DataLoader(valid_dataset, batch_size=self._batch_size, shuffle=True)
-
-        if self.model is not None and self.model.fc1.in_features != feature_dim:
+        if self.model is not None and self.model.in_features != feature_dim:
             logger.warning("The dimension of features doesn't match. A new model instance will be created.")
             self.model = None
         if self.model is None:
+            ffn_layers = kwargs.get('ffn_layers', 5)
+            ffn_expansion_factor = kwargs.get('ffn_expansion_factor', 2)
+            ffn_bias = kwargs.get('ffn_bias', True)
+            ffn_dropout_rate = kwargs.get('ffn_dropout_rate', 0.1)
             self.model = LongContextRecursiveSequential(input_dim=feature_dim, output_dim=1,
                                                         hidden_dim=hidden_dim,
                                                         recursive_layers=recursive_layers,
+                                                        ffn_layers=ffn_layers,
+                                                        ffn_expansion_factor=ffn_expansion_factor,
+                                                        ffn_bias=ffn_bias,
+                                                        ffn_dropout_rate=ffn_dropout_rate,
+                                                        attn_proj_layers=kwargs.get('attn_proj_layers', ffn_layers),
+                                                        attn_proj_bias=kwargs.get('attn_proj_bias', ffn_bias),
+                                                        attn_proj_expansion_factor=kwargs.get('attn_proj_expansion_factor', ffn_expansion_factor),
+                                                        attn_proj_dropout_rate=kwargs.get('attn_proj_dropout_rate', ffn_dropout_rate),
                                                         device=self.device, dtype=dtype)
+        logger.debug(f'Training Model: {self.model}')
         loss_function = nn.BCELoss()
         optimizer = optim.Adamax(self.model.parameters(), lr=learning_rate)
         for epoch in range(num_epochs):

{deeplotx-0.5.5.dist-info → deeplotx-0.6.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: deeplotx
-Version: 0.5.5
+Version: 0.6.1
 Summary: Easy-2-use long text NLP toolkit.
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
@@ -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:
@@ -265,7 +288,8 @@ Dynamic: license-file
     long_text_encoder = LongTextEncoder(
         max_length=2048,  # 最大文本大小, 超出截断
         chunk_size=448,  # 块大小 (按 Token 计)
-        overlapping=32  # 块间重叠大小 (按 Token 计)
+        overlapping=32,  # 块间重叠大小 (按 Token 计)
+        cache_capacity=512  # 缓存大小
     )
 
     trainer = TextBinaryClassifierTrainer(

deeplotx-0.6.1.dist-info/RECORD

@@ -0,0 +1,31 @@
+deeplotx/__init__.py,sha256=S0hLmRkHdoaxv7IPKVXh5Oat27pt_FGKGmKjp3aAyMU,1129
+deeplotx/encoder/__init__.py,sha256=BrsF5_4O-4pfihYF2wjExDOoAY-03kGJTH-Mhez4tsE,129
+deeplotx/encoder/encoder.py,sha256=oSBdA-MiwMKNfTFJWR-RdvNS0G0qfX-Qchwy4LuwB00,3985
+deeplotx/encoder/long_text_encoder.py,sha256=It0hXuSe0Hq5Y_3QhjEqvF1JbtX6Hc2VzVabzOu7fLA,3625
+deeplotx/encoder/longformer_encoder.py,sha256=A8FXqd4mdHxSn_o_R689XtpT73ISDT788EgMQRGLC2g,1822
+deeplotx/nn/__init__.py,sha256=f7f6Qx1Xkw3Nn3Lvafe7Pq4pUO7ZcESIA8KZxnSL_OM,535
+deeplotx/nn/auto_regression.py,sha256=8eEdXhOjRLKP4MpgX1wt9L1grU4_fS49ejVoNzFs7LM,877
+deeplotx/nn/base_neural_network.py,sha256=s7jHe7HprOelD1wZRbFdqb5Hxqs3sjLXLEo7OyDBHtk,3215
+deeplotx/nn/feed_forward.py,sha256=3lWV_snCp_PiqjxTYoiNlL9EF2heekWbMkKXoPlljkM,2839
+deeplotx/nn/linear_regression.py,sha256=QybSRfsf9PpgJAWixvrSNn3OYRKJXpSZMfqdzpw-Kd8,1280
+deeplotx/nn/logistic_regression.py,sha256=WfgHVNGIvAYsX2iea2wRlLgfbubYWyZkBLYpnpwOiyU,937
+deeplotx/nn/long_context_auto_regression.py,sha256=oMrxeVuCa1M2EQJSbOYlpTjl5NrkKGAHers8qIaZdU8,911
+deeplotx/nn/long_context_recursive_sequential.py,sha256=sU_22QH7Z6EJurMbTVEYPd83wC2dzadMIeztVIcc04I,2173
+deeplotx/nn/recursive_sequential.py,sha256=WsmXaIgTdpudo2bYcpBX8bKeJgPnT-atwEmLSXqQEco,2743
+deeplotx/nn/self_attention.py,sha256=HW9ZB3S6-yfTQc2745rJ6TM7L01P8ewxt7nGHosE2r8,2291
+deeplotx/nn/softmax_regression.py,sha256=PN_1Zr_B_z5zYC_s_8k6c5fllOtxfJEvVvCmC9GRmx0,958
+deeplotx/similarity/__init__.py,sha256=s3u-KSgxjnMcWpIItKgXNltFMPQ7YY3CqsqHI-5F1c8,724
+deeplotx/similarity/distribution.py,sha256=wQGouuuW531pZeBRKBujXsdsoz4fDnPw7_GW81jwepc,1066
+deeplotx/similarity/set.py,sha256=zhGFxtSIXlWqvipBYzoiPahp4g0boAIoUiMfG0wl07A,686
+deeplotx/similarity/vector.py,sha256=WVbDHqykt-fvuILVrhUCtIFAOEjY_zvttrXGM9eylG0,1125
+deeplotx/trainer/__init__.py,sha256=Fl5DR9UecQc5VtBcczU9sx_HtPNoFohpuELOh-Jrsks,77
+deeplotx/trainer/base_trainer.py,sha256=z0MeAT-rRYmjeBXt0ckt7J1itYArR0Cx02wHesXUoZE,385
+deeplotx/trainer/text_binary_classification_trainer.py,sha256=7oLzgXvdmFpQiBy7ncJ0smdqnMGr8xdZs6nTWpj6qfw,6085
+deeplotx/util/__init__.py,sha256=JxqAK_WOOHcYVSTHBT1-WuBwWrPEVDTV3titeVWvNUM,74
+deeplotx/util/hash.py,sha256=qbNU3RLBWGQYFVte9WZBAkZ1BkdjCXiKLDaKPN54KFk,662
+deeplotx/util/read_file.py,sha256=ptzouvEQeeW8KU5BrWNJlXw-vFXVrpS9SkAUxsu6A8A,612
+deeplotx-0.6.1.dist-info/licenses/LICENSE,sha256=IwGE9guuL-ryRPEKi6wFPI_zOhg7zDZbTYuHbSt_SAk,35823
+deeplotx-0.6.1.dist-info/METADATA,sha256=a1KcBHaewfyOwIywZ3wtBr8mdly4ofdb7Z4g2KYVzUk,12251
+deeplotx-0.6.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+deeplotx-0.6.1.dist-info/top_level.txt,sha256=hKg4pVDXZ-WWxkRfJFczRIll1Sv7VyfKCmzHLXbuh1U,9
+deeplotx-0.6.1.dist-info/RECORD,,

deeplotx-0.5.5.dist-info/RECORD

@@ -1,30 +0,0 @@
-deeplotx/__init__.py,sha256=6El66QXHDrgNMsNIG9bG97WO8BhPK5btXbTikzx2ce4,1087
-deeplotx/encoder/__init__.py,sha256=BrsF5_4O-4pfihYF2wjExDOoAY-03kGJTH-Mhez4tsE,129
-deeplotx/encoder/encoder.py,sha256=p1e4Dx3-Ghdl0MGNalr0D_OnafwaJnbhscEDVq-y73A,2400
-deeplotx/encoder/long_text_encoder.py,sha256=GatkOF1QQHLtvyuikfCP4xpzfDvszJJyonaS9f7wSxg,3401
-deeplotx/encoder/longformer_encoder.py,sha256=A8FXqd4mdHxSn_o_R689XtpT73ISDT788EgMQRGLC2g,1822
-deeplotx/nn/__init__.py,sha256=CS0UwyYKa8wI6vu6FBIYxvm-HAmw39MTMFlZDtqi6UA,444
-deeplotx/nn/auto_regression.py,sha256=7P63opWCWMqE2DigwbsL6kfXtFtJPz00Yo1RqflBz4A,572
-deeplotx/nn/base_neural_network.py,sha256=o9s0NqxkDcFZdipX8UrlbBmwYHOg7wPmzbjBEeGw63s,2902
-deeplotx/nn/linear_regression.py,sha256=7TbbplBgY70b1l5lKvTJMzDWQ8khQfnRCyMjObhVdEc,2180
-deeplotx/nn/logistic_regression.py,sha256=YiSLAon8gLDtMXAkPQ210sauod24eyJYYH50fPhj6T8,667
-deeplotx/nn/long_context_auto_regression.py,sha256=Z67Enq1kc1bERIrQW4jHeDQQmisOXhhjrtaPklnHkyw,605
-deeplotx/nn/long_context_recursive_sequential.py,sha256=_fKpPA7wt6B0kPyyig4xuhmLxygK19FSLgxW1Xa453M,1487
-deeplotx/nn/recursive_sequential.py,sha256=8YHZ-IdLyMJN5QVWPMuizDxLodAE9Bgdg1_YtIxFw7o,2247
-deeplotx/nn/self_attention.py,sha256=fb34wXnfgAGYJEhqa1l9AxMa-AHcCTOLbUlAfaGIK7Q,1766
-deeplotx/nn/softmax_regression.py,sha256=BeVk0G2H3zKG6bsQgPRNWuTxnnNmVI2zFZtCHgARAAc,688
-deeplotx/similarity/__init__.py,sha256=s3u-KSgxjnMcWpIItKgXNltFMPQ7YY3CqsqHI-5F1c8,724
-deeplotx/similarity/distribution.py,sha256=wQGouuuW531pZeBRKBujXsdsoz4fDnPw7_GW81jwepc,1066
-deeplotx/similarity/set.py,sha256=zhGFxtSIXlWqvipBYzoiPahp4g0boAIoUiMfG0wl07A,686
-deeplotx/similarity/vector.py,sha256=WVbDHqykt-fvuILVrhUCtIFAOEjY_zvttrXGM9eylG0,1125
-deeplotx/trainer/__init__.py,sha256=Fl5DR9UecQc5VtBcczU9sx_HtPNoFohpuELOh-Jrsks,77
-deeplotx/trainer/base_trainer.py,sha256=z0MeAT-rRYmjeBXt0ckt7J1itYArR0Cx02wHesXUoZE,385
-deeplotx/trainer/text_binary_classification_trainer.py,sha256=BNBQdpaD8nB1dQv8naHNIravNcQC8JjOMqD-WRSrUH0,4931
-deeplotx/util/__init__.py,sha256=JxqAK_WOOHcYVSTHBT1-WuBwWrPEVDTV3titeVWvNUM,74
-deeplotx/util/hash.py,sha256=qbNU3RLBWGQYFVte9WZBAkZ1BkdjCXiKLDaKPN54KFk,662
-deeplotx/util/read_file.py,sha256=ptzouvEQeeW8KU5BrWNJlXw-vFXVrpS9SkAUxsu6A8A,612
-deeplotx-0.5.5.dist-info/licenses/LICENSE,sha256=IwGE9guuL-ryRPEKi6wFPI_zOhg7zDZbTYuHbSt_SAk,35823
-deeplotx-0.5.5.dist-info/METADATA,sha256=QE1R1jodTrnPFY7cbu4mQNPt8_BgKNJuHoSDswopueo,10880
-deeplotx-0.5.5.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-deeplotx-0.5.5.dist-info/top_level.txt,sha256=hKg4pVDXZ-WWxkRfJFczRIll1Sv7VyfKCmzHLXbuh1U,9
-deeplotx-0.5.5.dist-info/RECORD,,