deeplotx 0.4.10__tar.gz → 0.4.12b0__tar.gz

{deeplotx-0.4.10 → deeplotx-0.4.12b0}/PKG-INFO

@@ -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 → deeplotx-0.4.12b0}/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-0.4.10 → deeplotx-0.4.12b0}/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-0.4.10 → deeplotx-0.4.12b0}/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-0.4.10 → deeplotx-0.4.12b0}/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 → deeplotx-0.4.12b0}/deeplotx.egg-info/PKG-INFO

@@ -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 → deeplotx-0.4.12b0}/deeplotx.egg-info/SOURCES.txt

@@ -27,4 +27,5 @@ deeplotx/trainer/base_trainer.py
 deeplotx/trainer/text_binary_classification_trainer.py
 deeplotx/util/__init__.py
 deeplotx/util/hash.py
-deeplotx/util/read_file.py
+deeplotx/util/read_file.py
+test/test.py

{deeplotx-0.4.10 → deeplotx-0.4.12b0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "deeplotx"
-version = "0.4.10"
+version = "0.4.12.beta"
 description = "Easy-2-use long text NLP toolkit."
 readme = "README.md"
 requires-python = ">=3.10"

deeplotx-0.4.12b0/test/test.py

@@ -0,0 +1,21 @@
+from deeplotx import TextBinaryClassifierTrainer, LongTextEncoder
+from deeplotx.util import get_files, read_file
+
+long_text_encoder = LongTextEncoder(
+    max_length=2048,
+    chunk_size=512,
+    overlapping=128
+)
+
+trainer = TextBinaryClassifierTrainer(
+    long_text_encoder=long_text_encoder,
+    batch_size=4,
+    train_ratio=0.9
+)
+
+pos_data_path = './data/pos'
+neg_data_path = './data/neg'
+pos_data = [read_file(x) for x in get_files(pos_data_path)]
+neg_data = [read_file(x) for x in get_files(neg_data_path)]
+model = trainer.train(pos_data, neg_data, num_epochs=20, learning_rate=2e-5, train_loss_threshold=1)
+model.predict(long_text_encoder.encode('这是一个测试文本.').squeeze())