deeplotx 0.8.8__tar.gz → 0.9.0__tar.gz

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

@@ -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 → deeplotx-0.9.0}/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-0.9.0/deeplotx/ner/__init__.py

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

deeplotx-0.9.0/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-0.9.0/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-0.9.0/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 → deeplotx-0.9.0}/deeplotx.egg-info/PKG-INFO

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

@@ -11,6 +11,10 @@ deeplotx/encoder/__init__.py
 deeplotx/encoder/encoder.py
 deeplotx/encoder/long_text_encoder.py
 deeplotx/encoder/longformer_encoder.py
+deeplotx/ner/__init__.py
+deeplotx/ner/base_ner.py
+deeplotx/ner/bert_ner.py
+deeplotx/ner/named_entity.py
 deeplotx/nn/__init__.py
 deeplotx/nn/attention.py
 deeplotx/nn/auto_regression.py

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

@@ -3,8 +3,10 @@ jupyter
 numpy
 protobuf
 python-dotenv
+sentencepiece
 tiktoken
 torch
 transformers
 typing-extensions
 vortezwohl>=0.0.8
+name2gender>=0.0.4a0

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

@@ -1,6 +1,6 @@
 [project]
 name = "deeplotx"
-version = "0.8.8"
+version = "0.9.0"
 description = "Easy-2-use long text NLP toolkit."
 readme = "README.md"
 requires-python = ">=3.10"
@@ -10,9 +10,11 @@ dependencies = [
     "numpy",
     "protobuf",
     "python-dotenv",
+    "sentencepiece",
     "tiktoken",
     "torch",
     "transformers",
     "typing-extensions",
     "vortezwohl>=0.0.8",
+    "name2gender>=0.0.4a0",
 ]