deeplotx 0.8.6__tar.gz → 0.8.8__tar.gz

{deeplotx-0.8.6 → deeplotx-0.8.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: deeplotx
-Version: 0.8.6
+Version: 0.8.8
 Summary: Easy-2-use long text NLP toolkit.
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
@@ -10,6 +10,7 @@ Requires-Dist: jupyter
 Requires-Dist: numpy
 Requires-Dist: protobuf
 Requires-Dist: python-dotenv
+Requires-Dist: tiktoken
 Requires-Dist: torch
 Requires-Dist: transformers
 Requires-Dist: typing-extensions

{deeplotx-0.8.6 → deeplotx-0.8.8}/deeplotx/encoder/encoder.py

@@ -43,9 +43,11 @@ class Encoder(nn.Module):
         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 forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor, cls_only: bool = True,
+                *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, :]
+            emb_seq = self.encoder.forward(_input_tup[0], attention_mask=_input_tup[1]).last_hidden_state
+            return emb_seq[:, 0, :] if cls_only else emb_seq
 
         num_chunks = math.ceil(input_ids.shape[-1] / self.embed_dim)
         chunks, chunk_results = [], []
@@ -58,9 +60,9 @@ class Encoder(nn.Module):
         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)
+        return torch.cat(chunk_results, dim=-1) if cls_only else torch.cat(chunk_results, dim=-2)
 
-    def encode(self, text: str) -> torch.Tensor:
+    def encode(self, text: str, cls_only: bool = True) -> 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()
+        return self.forward(input_ids=_input_ids, attention_mask=_att_mask, cls_only=cls_only).squeeze()

{deeplotx-0.8.6 → deeplotx-0.8.8}/deeplotx/encoder/long_text_encoder.py

@@ -25,7 +25,7 @@ class LongTextEncoder(Encoder):
         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)
+        return idx, super().forward(x, attention_mask=mask, cls_only=True)
 
     @override
     def forward(self, text: str, flatten: bool = False, *args, **kwargs) -> torch.Tensor:

{deeplotx-0.8.6 → deeplotx-0.8.8}/deeplotx/encoder/longformer_encoder.py

@@ -9,7 +9,7 @@ from requests.exceptions import ConnectTimeout, SSLError
 from deeplotx import __ROOT__
 
 CACHE_PATH = os.path.join(__ROOT__, '.cache')
-DEFAULT_LONGFORMER = 'allenai/longformer-base-4096'
+DEFAULT_LONGFORMER = 'severinsimmler/xlm-roberta-longformer-base-16384'
 logger = logging.getLogger('deeplotx.embedding')
 
 
@@ -41,15 +41,16 @@ class LongformerEncoder(nn.Module):
                                                      trust_remote_code=True, local_files_only=True).to(self.device)
         logger.debug(f'{LongformerEncoder.__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, cls_only: bool = True) -> torch.Tensor:
         ori_mode = self.encoder.training
         self.encoder.eval()
         with torch.no_grad():
-            res = self.encoder.forward(input_ids, attention_mask=attention_mask).last_hidden_state[:, 0, :]
+            emb_seq = self.encoder.forward(input_ids, attention_mask=attention_mask).last_hidden_state
+            res = emb_seq[:, 0, :] if cls_only else emb_seq
         self.encoder.train(mode=ori_mode)
         return res
 
-    def encode(self, text: str) -> torch.Tensor:
+    def encode(self, text: str, cls_only: bool = True) -> 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()
+        return self.forward(input_ids=_input_ids, attention_mask=_att_mask, cls_only=cls_only).squeeze()

{deeplotx-0.8.6 → deeplotx-0.8.8}/deeplotx/nn/base_neural_network.py

@@ -99,11 +99,14 @@ class BaseNeuralNetwork(nn.Module):
 
     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
+        training_state_dict = dict()
+        for m in self.modules():
+            training_state_dict[m] = m.training
+            m.training = False
         with torch.no_grad():
             res = self.forward(x)
-        self.training = __train
+        for m, training_state in training_state_dict.items():
+            m.training = training_state
         return res
 
     def save(self, model_name: str | None = None, model_dir: str = '.', _suffix: str = DEFAULT_SUFFIX):

{deeplotx-0.8.6 → deeplotx-0.8.8}/deeplotx/nn/feed_forward.py

@@ -28,7 +28,7 @@ class FeedForwardUnit(BaseNeuralNetwork):
         x = self.layer_norm(x)
         x = self.up_proj(x)
         x = self.parametric_relu(x)
-        if self._dropout_rate > .0:
+        if self._dropout_rate > .0 and self.training:
             x = torch.dropout(x, p=self._dropout_rate, train=self.training)
         return self.down_proj(x) + residual
 

{deeplotx-0.8.6 → deeplotx-0.8.8}/deeplotx/nn/recursive_sequential.py

@@ -41,9 +41,12 @@ class RecursiveSequential(BaseNeuralNetwork):
 
     @override
     def predict(self, x: torch.Tensor) -> torch.Tensor:
-        __train = self.training
-        self.training = False
+        training_state_dict = dict()
+        for m in self.modules():
+            training_state_dict[m] = m.training
+            m.training = False
         with torch.no_grad():
             res = self.forward(x.unsqueeze(0), self.initial_state(batch_size=1))[0]
-        self.training = __train
+        for m, training_state in training_state_dict.items():
+            m.training = training_state
         return res

{deeplotx-0.8.6 → deeplotx-0.8.8}/deeplotx.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: deeplotx
-Version: 0.8.6
+Version: 0.8.8
 Summary: Easy-2-use long text NLP toolkit.
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
@@ -10,6 +10,7 @@ Requires-Dist: jupyter
 Requires-Dist: numpy
 Requires-Dist: protobuf
 Requires-Dist: python-dotenv
+Requires-Dist: tiktoken
 Requires-Dist: torch
 Requires-Dist: transformers
 Requires-Dist: typing-extensions

{deeplotx-0.8.6 → deeplotx-0.8.8}/deeplotx.egg-info/requires.txt

@@ -3,6 +3,7 @@ jupyter
 numpy
 protobuf
 python-dotenv
+tiktoken
 torch
 transformers
 typing-extensions

{deeplotx-0.8.6 → deeplotx-0.8.8}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "deeplotx"
-version = "0.8.6"
+version = "0.8.8"
 description = "Easy-2-use long text NLP toolkit."
 readme = "README.md"
 requires-python = ">=3.10"
@@ -10,6 +10,7 @@ dependencies = [
     "numpy",
     "protobuf",
     "python-dotenv",
+    "tiktoken",
     "torch",
     "transformers",
     "typing-extensions",