deeplotx 0.4.10__py3-none-any.whl → 0.4.12b0__py3-none-any.whl

deeplotx/encoder/bert_encoder.py

@@ -16,7 +16,8 @@ logger = logging.getLogger('deeplotx.embedding')
 class BertEncoder(nn.Module):
     def __init__(self, model_name_or_path: str = DEFAULT_BERT, device: str | None = None):
         super().__init__()
-        self.device = device if device is not None else torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        self.device = torch.device(device) if device is not None \
+            else torch.device('cuda' if torch.cuda.is_available() else 'cpu')
         self.tokenizer = BertTokenizer.from_pretrained(pretrained_model_name_or_path=model_name_or_path,
                                                        cache_dir=CACHE_PATH, _from_auto=True)
         self.bert = BertModel.from_pretrained(pretrained_model_name_or_path=model_name_or_path,

deeplotx/encoder/long_text_encoder.py

@@ -28,7 +28,7 @@ class LongTextEncoder(BertEncoder):
         def postprocess(tensors: list[torch.Tensor], _flatten: bool) -> torch.Tensor:
             if not _flatten:
                 return torch.stack(tensors, dim=0).squeeze()
-            _fin_emb_tensor = torch.tensor([], dtype=tensors[0].dtype)
+            _fin_emb_tensor = torch.tensor([], dtype=tensors[0].dtype, device=self.device)
             for _emb in tensors:
                 _fin_emb_tensor = torch.cat((_fin_emb_tensor.detach().clone(), _emb.detach().clone()), dim=-1)
             return _fin_emb_tensor.squeeze()
@@ -55,8 +55,8 @@ class LongTextEncoder(BertEncoder):
         for i in range(num_chunks):
             _tmp_left = max(i * self._chunk_size - self._overlapping, 0)
             _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),
-                           torch.tensor([_text_to_input_ids_att_mask[_tmp_left: _tmp_right]], dtype=torch.int)))
+            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)))
         with ThreadPoolExecutor(max_workers=min(num_chunks + 1, 3)) as executor:
             embeddings = list(executor.map(self.__chunk_embedding, chunks))
         embeddings.sort(key=lambda x: x[0])

deeplotx/encoder/longformer_encoder.py

@@ -15,7 +15,8 @@ logger = logging.getLogger('deeplotx.embedding')
 class LongformerEncoder(nn.Module):
     def __init__(self, model_name_or_path: str = DEFAULT_LONGFORMER, device: str | None = None):
         super().__init__()
-        self.device = device if device is not None else torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        self.device = torch.device(device) if device is not None \
+            else torch.device('cuda' if torch.cuda.is_available() else 'cpu')
         self.tokenizer = LongformerTokenizer.from_pretrained(pretrained_model_name_or_path=model_name_or_path,
                                                              cache_dir=CACHE_PATH, _from_auto=True)
         self.bert = LongformerModel.from_pretrained(pretrained_model_name_or_path=model_name_or_path,

deeplotx/trainer/text_binary_classification_trainer.py

@@ -16,6 +16,7 @@ class TextBinaryClassifierTrainer(BaseTrainer):
     def __init__(self, long_text_encoder: LongTextEncoder, batch_size: int = 2, train_ratio: float = 0.8):
         super().__init__(batch_size=batch_size, train_ratio=train_ratio)
         self._long_text_encoder = long_text_encoder
+        self.device = self._long_text_encoder.device
 
     @override
     def train(self, positive_texts: list[str], negative_texts: list[str],
@@ -27,8 +28,8 @@ class TextBinaryClassifierTrainer(BaseTrainer):
             positive_texts = positive_texts[:min_length]
             negative_texts = negative_texts[:min_length]
         all_texts = positive_texts + negative_texts
-        labels = ([torch.tensor([1.0], dtype=torch.float32) for _ in range(len(positive_texts))]
-                  + [torch.tensor([0.0], dtype=torch.float32) for _ in range(len(negative_texts))])
+        labels = ([torch.tensor([1.0], dtype=torch.float32, device=self.device) for _ in range(len(positive_texts))]
+                  + [torch.tensor([0.0], dtype=torch.float32, device=self.device) for _ in range(len(negative_texts))])
         text_embeddings = [self._long_text_encoder.encode(x) for x in all_texts]
         feature_dim = text_embeddings[0].shape[-1]
         inputs = torch.stack(text_embeddings)
@@ -44,6 +45,7 @@ class TextBinaryClassifierTrainer(BaseTrainer):
             self.model = None
         if self.model is None:
             self.model = LogisticRegression(input_dim=feature_dim, output_dim=1)
+        self.model.to(self.device)
         loss_function = nn.BCELoss()
         optimizer = optim.Adamax(self.model.parameters(), lr=learning_rate)
         for epoch in range(num_epochs):

{deeplotx-0.4.10.dist-info → deeplotx-0.4.12b0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: deeplotx
-Version: 0.4.10
+Version: 0.4.12b0
 Summary: Easy-2-use long text NLP toolkit.
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown

{deeplotx-0.4.10.dist-info → deeplotx-0.4.12b0.dist-info}/RECORD

@@ -1,8 +1,8 @@
 deeplotx/__init__.py,sha256=wMN_AI14V-0BPbQghYpvd2y7eUGfhr7jKTTuur-5Upg,1002
 deeplotx/encoder/__init__.py,sha256=EM-xrTsHoGaiiFpj-iFAxilMHXC_sQKWYrcq1qCnI3U,138
-deeplotx/encoder/bert_encoder.py,sha256=VCmYsBSqB9bRL_ge4bYssyx-Xy4oR0-DE1cMTuTn1tU,2412
-deeplotx/encoder/long_text_encoder.py,sha256=7On6NuaINDZLqgb3HsSJBEzbWXNZPh_MXAvO5KY471k,3313
-deeplotx/encoder/longformer_encoder.py,sha256=J8Si8Ta0bh7Vo7YsV0XdC7jGrvIt54GKcHr_pq2qHbI,1857
+deeplotx/encoder/bert_encoder.py,sha256=IZsmkcmK6ulwTS4zubljW5uvq2r1Ik03nYG6jqcgQL8,2441
+deeplotx/encoder/long_text_encoder.py,sha256=hl_O8kR9o1kcII9YfSx2rf_Pk0l_Rv7LNbsS9UsTU0c,3373
+deeplotx/encoder/longformer_encoder.py,sha256=vsDIiS9kLzvAalDnCGfTuAd2gfqDBgKUqPj6tPgF-BM,1886
 deeplotx/nn/__init__.py,sha256=oQ-vYXyuaGelfCOs2im_gZXAiiBlCCVXh1uw9yjvRMs,253
 deeplotx/nn/auto_regression.py,sha256=o82C9TREZbhGdj2knSVGTXhjJne0LGEqc7BllByJJWE,449
 deeplotx/nn/base_neural_network.py,sha256=xWKG4FX6Jzdlrfc1HOW1aO9uh0Af3D-dB5Jl7eCxsAk,1635
@@ -16,12 +16,12 @@ deeplotx/similarity/set.py,sha256=zhGFxtSIXlWqvipBYzoiPahp4g0boAIoUiMfG0wl07A,68
 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=5O-5dwVMCj5EDX9gjJwCA468OR4UozJ7V8b-JxeUB0s,4080
+deeplotx/trainer/text_binary_classification_trainer.py,sha256=NhLFndk4I1PViSfat4XadDV-vSUGPOZ0RabfhQ5FLKY,4210
 deeplotx/util/__init__.py,sha256=JxqAK_WOOHcYVSTHBT1-WuBwWrPEVDTV3titeVWvNUM,74
 deeplotx/util/hash.py,sha256=wwsC6kOQvbpuvwKsNQOARd78_wePmW9i3oaUuXRUnpc,352
 deeplotx/util/read_file.py,sha256=ptzouvEQeeW8KU5BrWNJlXw-vFXVrpS9SkAUxsu6A8A,612
-deeplotx-0.4.10.dist-info/licenses/LICENSE,sha256=IwGE9guuL-ryRPEKi6wFPI_zOhg7zDZbTYuHbSt_SAk,35823
-deeplotx-0.4.10.dist-info/METADATA,sha256=2x4T_toVpNHl8eXdLmHwKjjjvOPEsTB2qVg3pvyLltA,6285
-deeplotx-0.4.10.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-deeplotx-0.4.10.dist-info/top_level.txt,sha256=hKg4pVDXZ-WWxkRfJFczRIll1Sv7VyfKCmzHLXbuh1U,9
-deeplotx-0.4.10.dist-info/RECORD,,
+deeplotx-0.4.12b0.dist-info/licenses/LICENSE,sha256=IwGE9guuL-ryRPEKi6wFPI_zOhg7zDZbTYuHbSt_SAk,35823
+deeplotx-0.4.12b0.dist-info/METADATA,sha256=2JQcCaNV4WQ2jYDw50dIT5aLcwoEBkkjrvDAOCrcdbM,6287
+deeplotx-0.4.12b0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+deeplotx-0.4.12b0.dist-info/top_level.txt,sha256=hKg4pVDXZ-WWxkRfJFczRIll1Sv7VyfKCmzHLXbuh1U,9
+deeplotx-0.4.12b0.dist-info/RECORD,,