deeplotx 0.8.8__py3-none-any.whl → 0.9.0__py3-none-any.whl

deeplotx/__init__.py

@@ -4,6 +4,7 @@ import os
 __ROOT__ = os.path.dirname(os.path.abspath(__file__))
 
 from .encoder import Encoder, LongTextEncoder, LongformerEncoder
+from .ner import BertNER, NamedEntity
 from .nn import (
     FeedForward,
     MultiHeadFeedForward,
@@ -40,3 +41,5 @@ logger = logging.getLogger('deeplotx.trainer')
 logger.setLevel(logging.DEBUG)
 logger = logging.getLogger('deeplotx.embedding')
 logger.setLevel(logging.DEBUG)
+logger = logging.getLogger('deeplotx.ner')
+logger.setLevel(logging.DEBUG)

deeplotx/ner/__init__.py

@@ -0,0 +1,3 @@
+from .named_entity import NamedEntity
+from .base_ner import BaseNER
+from .bert_ner import BertNER

deeplotx/ner/base_ner.py

@@ -0,0 +1,7 @@
+from deeplotx.ner.named_entity import NamedEntity
+
+
+class BaseNER:
+    def __init__(self): ...
+
+    def extract_entities(self, s: str, *args, **kwargs) -> list[NamedEntity]: ...

deeplotx/ner/bert_ner.py

@@ -0,0 +1,72 @@
+import logging
+import os
+from requests.exceptions import ConnectTimeout, SSLError
+
+import torch
+from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline
+
+from deeplotx import __ROOT__
+from deeplotx.ner.base_ner import BaseNER
+from deeplotx.ner.named_entity import NamedEntity
+
+CACHE_PATH = os.path.join(__ROOT__, '.cache')
+DEFAULT_BERT_NER = 'Davlan/xlm-roberta-base-ner-hrl'
+logger = logging.getLogger('deeplotx.ner')
+
+
+class BertNER(BaseNER):
+    def __init__(self, model_name_or_path: str = DEFAULT_BERT_NER, 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 = AutoModelForTokenClassification.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 = AutoModelForTokenClassification.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 = AutoModelForTokenClassification.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
+        self._ner_pipeline = pipeline(task='ner', model=self.encoder, tokenizer=self.tokenizer, trust_remote_code=True)
+        logger.debug(f'{BaseNER.__name__} initialized on device: {self.device}.')
+
+    def extract_entities(self, s: str, prob_threshold: float = .0, *args, **kwargs) -> list[NamedEntity]:
+        assert prob_threshold <= 1., f'prob_threshold ({prob_threshold}) cannot be larger than 1.'
+        s = ' ' + s
+        raw_entities = self._ner_pipeline(s)
+        entities = []
+        for ent in raw_entities:
+            entities.append([s[ent['start']: ent['end']], ent['entity'], ent['score'].item()])
+        while True:
+            for i, ent in enumerate(entities):
+                if len(ent[0].strip()) < 1:
+                    del entities[i]
+                if ent[1].upper().startswith('I') and entities[i - 1][1].upper().startswith('B'):
+                    entities[i - 1][0] += ent[0]
+                    entities[i - 1][2] *= ent[2]
+                    del entities[i]
+            _continue = False
+            for ent in entities:
+                if ent[1].upper().startswith('I'):
+                    _continue = True
+            if not _continue:
+                break
+        for ent in entities:
+            ent[0] = ent[0].strip()
+            if ent[1].upper().startswith('B'):
+                ent[1] = ent[1].upper()[1:].strip('-')
+        return [NamedEntity(*_) for _ in entities if _[2] >= prob_threshold]

deeplotx/ner/named_entity.py

@@ -0,0 +1,8 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class NamedEntity:
+    text: str
+    type: str
+    probability: float

{deeplotx-0.8.8.dist-info → deeplotx-0.9.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: deeplotx
-Version: 0.8.8
+Version: 0.9.0
 Summary: Easy-2-use long text NLP toolkit.
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
@@ -10,11 +10,13 @@ Requires-Dist: jupyter
 Requires-Dist: numpy
 Requires-Dist: protobuf
 Requires-Dist: python-dotenv
+Requires-Dist: sentencepiece
 Requires-Dist: tiktoken
 Requires-Dist: torch
 Requires-Dist: transformers
 Requires-Dist: typing-extensions
 Requires-Dist: vortezwohl>=0.0.8
+Requires-Dist: name2gender>=0.0.4a0
 Dynamic: license-file
 
 [![Ask DeepWiki](https://deepwiki.com/badge.svg)](https://deepwiki.com/vortezwohl/DeepLoTX)

{deeplotx-0.8.8.dist-info → deeplotx-0.9.0.dist-info}/RECORD

@@ -1,8 +1,12 @@
-deeplotx/__init__.py,sha256=xEq8WQ2LpEZoLX_Z464d0dy4aemFGrEV6ZMJr6ioFnQ,1186
+deeplotx/__init__.py,sha256=_CAKf1HvuHxUMoEYmYGaH8aSQVK99_LIXSzoF7vIviU,1301
 deeplotx/encoder/__init__.py,sha256=BrsF5_4O-4pfihYF2wjExDOoAY-03kGJTH-Mhez4tsE,129
 deeplotx/encoder/encoder.py,sha256=wVRl3p_7eg7qT_tJEit5qnmZx7dXkMVLxAtao5vImkk,4201
 deeplotx/encoder/long_text_encoder.py,sha256=4oRa9FqfGNZ8-gq14UKuhDkZC0A1Xi-wKmbQsn-uZ58,3966
 deeplotx/encoder/longformer_encoder.py,sha256=7Lm65AUD3qwbrzrhJ3dPZkyHeNRSapga3f-5QJCxV5A,3538
+deeplotx/ner/__init__.py,sha256=NhLSZ6Urvv9jokdZOKC5vQd0STKmm2gMrdcMgvjQULc,101
+deeplotx/ner/base_ner.py,sha256=OwD9fTuTu1xq2lC2NUnJ0MDJ3KGhUiBCXw9ySYg4eQQ,185
+deeplotx/ner/bert_ner.py,sha256=179cqfodG-eyOOD8s8QgnFprDQRBukuFByqLoMH4d6c,4420
+deeplotx/ner/named_entity.py,sha256=NfwdnLntkMrfPeZnIuWYbVnAvf2sdKQU1naodqRy0b0,125
 deeplotx/nn/__init__.py,sha256=YILwbxb-NHdiJjfOwBKH8F7PuZSDZSrGpTznPDucTro,710
 deeplotx/nn/attention.py,sha256=R-i-Rd7gnsh6hwXDeYfqLQOJvfSZIGfQbFzRlC91XLo,2879
 deeplotx/nn/auto_regression.py,sha256=j_R7WGPq9REngjpLuX5c0AaNqOpgGm2Vfrolw-XjWXw,877
@@ -28,8 +32,8 @@ deeplotx/trainer/text_binary_classification_trainer.py,sha256=TFxOX8rWU_zKliI9zm
 deeplotx/util/__init__.py,sha256=5CH4MTeSgsmCe3LPMfvKoSBpwh6jDSBuHVElJvzQzgs,90
 deeplotx/util/hash.py,sha256=qbNU3RLBWGQYFVte9WZBAkZ1BkdjCXiKLDaKPN54KFk,662
 deeplotx/util/read_file.py,sha256=ptzouvEQeeW8KU5BrWNJlXw-vFXVrpS9SkAUxsu6A8A,612
-deeplotx-0.8.8.dist-info/licenses/LICENSE,sha256=IwGE9guuL-ryRPEKi6wFPI_zOhg7zDZbTYuHbSt_SAk,35823
-deeplotx-0.8.8.dist-info/METADATA,sha256=EhDWaoH6HmlnNga9c7VitZiBerZ2GAXXrz3BWG9latc,13163
-deeplotx-0.8.8.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-deeplotx-0.8.8.dist-info/top_level.txt,sha256=hKg4pVDXZ-WWxkRfJFczRIll1Sv7VyfKCmzHLXbuh1U,9
-deeplotx-0.8.8.dist-info/RECORD,,
+deeplotx-0.9.0.dist-info/licenses/LICENSE,sha256=IwGE9guuL-ryRPEKi6wFPI_zOhg7zDZbTYuHbSt_SAk,35823
+deeplotx-0.9.0.dist-info/METADATA,sha256=9dx6Tz81e9YWTnsy-1VPw30ts9oJUcQD0Kohm36_HKA,13230
+deeplotx-0.9.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+deeplotx-0.9.0.dist-info/top_level.txt,sha256=hKg4pVDXZ-WWxkRfJFczRIll1Sv7VyfKCmzHLXbuh1U,9
+deeplotx-0.9.0.dist-info/RECORD,,