streamlit-octostar-utils 0.1.7a4__tar.gz → 0.1.7a6__tar.gz

{streamlit_octostar_utils-0.1.7a4 → streamlit_octostar_utils-0.1.7a6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: streamlit-octostar-utils
-Version: 0.1.7a4
+Version: 0.1.7a6
 Summary: 
 License: MIT
 Author: Octostar

{streamlit_octostar_utils-0.1.7a4 → streamlit_octostar_utils-0.1.7a6}/pyproject.toml

@@ -5,7 +5,7 @@ include = '\.pyi?$'
 
 [tool.poetry]
 name = "streamlit-octostar-utils"
-version = "0.1.7a4"
+version = "0.1.7a6"
 description = ""
 license = "MIT"
 authors = ["Octostar"]

{streamlit_octostar_utils-0.1.7a4 → streamlit_octostar_utils-0.1.7a6}/streamlit_octostar_utils/nlp/ner.py

@@ -14,6 +14,7 @@ import numpy as np
 import math
 import nltk
 from typing import Optional, List
+from pydantic import BaseModel, ConfigDict, Field
 
 SPACY_NER_MODELS = {
     "english": lambda: load_spacy(
@@ -24,15 +25,15 @@ SPACY_NER_MODELS = {
 FLAIR_NER_MODELS = {"english": lambda: SequenceTagger.load("flair/ner-english")}
 REGEX_NER_MODELS = {
     "IP_ADDRESS": [
-        r"(?:(?<=:=)|(?<=\s)|(?<=\b))(?:\d{1,3}\.){3}\d{1,3}(?::\d{1,5})?(?:(?=\s)|(?=\b))",
-        r"(?:(?<=:=)|(?<=\s)|(?<=\b))(?:[A-Fa-f0-9]{1,4}:){7}[A-Fa-f0-9]{1,4}(?::\d{1,5})?(?:(?=\s)|(?=\b))"
+        r"\b(?:(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.){3}(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)(?::(?:[0-9]|[1-9][0-9]{1,3}|[1-5][0-9]{4}|6[0-4][0-9]{3}|65[0-4][0-9]{2}|655[0-2][0-9]|6553[0-5]))?\b",
     ],
-    "PHONE": r"(?:(?<=:=)|(?<=\s)|(?<=\b))[+]?[(]?[0-9]{1,4}[)]?[-\s\/0-9]*(?:(?=\s)|(?=\b))",
-    "EMAIL": r"(?:(?<=:=)|(?<=\s)|(?<=\b))[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}(?:(?=\s)|(?=\b))",
+    "PHONE": r"(?:(?:\+(?:\d{1,3}[ .-]?)?(?:\(\d{1,3}\)[ .-]?)?)(?:\d{2,5}[ .-]?){1,3}|\d{2,5}[ .-]\d{2,5}(?:[ .-]\d{2,5}){0,2})\b",
+    "EMAIL": r"\b[a-zA-Z0-9._%+-]+@[a-zA-Z0-9](?:[a-zA-Z0-9-]*[a-zA-Z0-9])?(?:\.[a-zA-Z0-9](?:[a-zA-Z0-9-]*[a-zA-Z0-9])?)+\b",
+    "URL": r"\b(?:(?:https?|ftp|sftp|ftps|ssh|file|mailto|git|onion|ipfs|ipns):\/\/|www\.)(?:[a-z0-9](?:[a-z0-9-]{0,61}[a-z0-9])?\.)+[a-z]{2,}(?::\d+)?(?:\/(?:[-a-z0-9\/_.,~%+:@]|(?:%[0-9a-f]{2}))*)?(?:\?(?:[-a-z0-9\/_.,~%+:@=&]|(?:%[0-9a-f]{2}))*)?(?:#(?:[-a-z0-9\/_.,~%+:@=&]|(?:%[0-9a-f]{2}))*)?|(?:https?:\/\/)?[a-z2-7]{16,56}\.onion(?:\/(?:[-a-z0-9\/_.,~%+:@]|(?:%[0-9a-f]{2}))*)?(?:\?(?:[-a-z0-9\/_.,~%+:@=&]|(?:%[0-9a-f]{2}))*)?(?:#(?:[-a-z0-9\/_.,~%+:@=&]|(?:%[0-9a-f]{2}))*)\b",
 }
 
 BASE_TO_ONTONOTES_LABELMAP = {"PER": "PERSON"}
-BASE_ALLOWED_LABELS = ["PERSON", "ORG", "LOC", "NORP", "GPE", "PRODUCT", "DATE", "TIME", "PHONE"]
+BASE_ALLOWED_LABELS = ["PERSON", "ORG", "LOC", "NORP", "GPE", "PRODUCT", "DATE", "TIME", "PHONE", "IP_ADDRESS", "EMAIL"]
 
 
 def _sumy__get_best_sentences(sentences, rating, *args, **kwargs):
@@ -75,33 +76,21 @@ def get_nltk_tokenizer(language: str) -> Tokenizer:
     return Tokenizer(language)
 
 
-class NERObject(object):
-    def __init__(self, name, label, score, context, count, comentions):
-        self.name: str = name
-        self.label: str = label
-        self.score: Optional[float] = score
-        self.context: Optional[str] = context
-        self.count: int = count
-        self.comentions: Optional[List[str]] = comentions
-        self.sources: Optional[List[str]] = list()
+class NERObject(BaseModel):
+    name: str
+    label: str
+    score: float = 0.0
+    start: int
+    count: int
+    context: str | None = None
+    comentions: list[str] = Field(default_factory=list)
+    model_config = ConfigDict(extra="allow")
 
-    def to_dict(self):
-        data = {
-            "name": self.name,
-            "label": self.label,
-            "score": self.score,
-            "context": self.context,
-            "count": self.count,
-            "comentions": self.comentions or [],
-        }
-        if self.sources:
-            data["sources"] = self.sources
-            
     def __repr__(self):
         return f"NERObject(label={self.label},name={self.name})"
 
 
-def postprocess_ner(entities, whitelisted_labels=None, max_entities=None):
+def postprocess_ner(entities: list[NERObject], whitelisted_labels=None, max_entities=None):
     if whitelisted_labels is not None:
         entities = [e for e in entities if e.label in whitelisted_labels]
     entities = sorted(entities, key=lambda x: x.name)
@@ -109,29 +98,48 @@ def postprocess_ner(entities, whitelisted_labels=None, max_entities=None):
     for _, group in itertools.groupby(entities, key=lambda x: x.name):
         group = list(group)
         best_entity = max(group, key=lambda x: x.score * x.count)
-        best_entity = NERObject(
-            best_entity.name,
-            best_entity.label,
-            best_entity.score,
-            best_entity.context,
-            sum([0] + [e.count for e in group]),
-            list(set(itertools.chain(*[e.comentions for e in group]))),
-        )
-        best_entity.sources = list(set(itertools.chain(*[e.sources for e in group])))
-        final_entities.append(best_entity)
+        merged_data = {
+            "name": best_entity.name,
+            "label": best_entity.label,
+            "score": best_entity.score,
+            "context": best_entity.context,
+            "count": sum(e.count for e in group),
+            "start": best_entity.start,
+        }
+        all_fields = best_entity.model_fields.keys()
+        for field in all_fields:
+            if field in merged_data:
+                continue
+            values = [getattr(e, field, None) for e in group if getattr(e, field, None) is not None]
+            if not values:
+                continue
+            if isinstance(values[0], list):
+                merged_data[field] = list(set(itertools.chain.from_iterable(values or [])))
+            else:
+                merged_data[field] = getattr(best_entity, field, None)
+        final_entities.append(NERObject(**merged_data))
     final_entities = sorted(final_entities, key=lambda x: x.score * x.count, reverse=True)
     if max_entities and len(final_entities) > max_entities:
         final_entities = final_entities[:max_entities]
     return final_entities
 
 
-def compute_ner(language, sentences, spacy_model, flair_model=None, context_width=150):
+def compute_ner(
+    language,
+    sentences,
+    spacy_model,
+    flair_model=None,
+    context_width=150,
+    with_scores=True,
+    with_comentions=True,
+    with_context=True,
+):
     sentence_starts = [0] + [len(s[0]) + 1 for s in sentences]
     del sentence_starts[-1]
     sentence_starts = list(np.cumsum(sentence_starts))
     text = "\n".join([s[0] for s in sentences])
     min_score = 1.0
-    entities = []
+    entities: list[NERObject] = []
 
     # FLAIR model (if not fast)
     if flair_model:
@@ -139,88 +147,115 @@ def compute_ner(language, sentences, spacy_model, flair_model=None, context_widt
         flair_model.predict(input)
         output = [e for sentence in input for e in sentence.get_spans("ner")]
         flair_entities = [
-            (
-                entity.text,
-                BASE_TO_ONTONOTES_LABELMAP.get(
+            NERObject(
+                name=entity.text,
+                label=BASE_TO_ONTONOTES_LABELMAP.get(
                     entity.annotation_layers["ner"][0].value,
                     entity.annotation_layers["ner"][0].value,
                 ),
-                entity.score,
-                sentence_starts[input.index(entity[0].sentence)] + entity[0].start_position,
+                score=entity.score,
+                start=sentence_starts[input.index(entity[0].sentence)] + entity[0].start_position,
+                count=1,
             )
             for entity in output
         ]
-        min_score = min(min_score, *[e[2] for e in flair_entities])
+        min_score = min([min_score] + [e.score for e in flair_entities])
         entities += flair_entities
         del flair_entities
 
+    print("Checking REGEXES")
     # REGEX model
     for label, regexes in REGEX_NER_MODELS.items():
         if not isinstance(regexes, list):
             regexes = [regexes]
         for regex in regexes:
-            print(regex)
             regex_entities = [
-                (match.group(), label, min_score - 0.5, match.start()) for match in re.finditer(regex, text)
+                NERObject(
+                    name=match.group(),
+                    label=label,
+                    score=min_score - 0.5,
+                    count=1,
+                    start=match.start(),
+                )
+                for match in re.finditer(regex, text)
             ]
-            print(regex_entities)
             entities += regex_entities
-            
+    min_score = min([min_score] + [e.score for e in regex_entities])
+
     # SPACY model
-    spacy_entities = [
-        (
-            entity.text,
-            BASE_TO_ONTONOTES_LABELMAP.get(entity.label_, entity.label_),
-            min_score - 1,
-            entity.start_char,
-        )
-        for entity in spacy_model(text).ents
-    ]
-    entities += spacy_entities
-    del spacy_entities
+    print("CHECKING SPACY")
+    chunks = []
+    chunk_start_offsets = []
+    current_chunk = []
+    current_length = 0
+    offset = 0
+    for sentence, _ in sentences:
+        sentence_len = len(sentence) + 1
+        if sentence_len > spacy_model.max_length:
+            truncated = sentence[: spacy_model.max_length - 1]
+            chunks.append(truncated)
+            chunk_start_offsets.append(offset)
+            offset += sentence_len
+            continue
+        if current_length + sentence_len > spacy_model.max_length:
+            chunks.append("\n".join(current_chunk))
+            chunk_start_offsets.append(offset - current_length)
+            current_chunk = []
+            current_length = 0
+        current_chunk.append(sentence)
+        current_length += sentence_len
+        offset += sentence_len
+    if current_chunk:
+        chunks.append("\n".join(current_chunk))
+        chunk_start_offsets.append(offset - current_length)
+    for i, chunk in enumerate(chunks):
+        doc = spacy_model(chunk)
+        chunk_offset = chunk_start_offsets[i]
+        for entity in doc.ents:
+            entities.append(
+                NERObject(
+                    name=entity.text,
+                    label=BASE_TO_ONTONOTES_LABELMAP.get(entity.label_, entity.label_),
+                    score=min_score - 0.5,
+                    start=chunk_offset + entity.start_char,
+                    count=1,
+                )
+            )
 
     # Reformatting for consistency
-    if entities:
-        min_entity_score = min([e[2] for e in entities])
-        max_entity_score = max([min_entity_score] + [e[2] for e in entities])
+    if not entities:
+        return []
+    if with_scores:
+        min_entity_score = min([e.score for e in entities])
+        max_entity_score = max([e.score for e in entities])
         entity_score_range = 1 if min_entity_score == max_entity_score else (max_entity_score - min_entity_score)
-        entities = [(e[0], e[1], (e[2] - min_entity_score) / entity_score_range, e[3]) for e in entities]
-        scores = list(np.searchsorted(sentence_starts, [e[3] + 1 for e in entities]))
+        for e in entities:
+            e.score = (e.score - min_entity_score) / entity_score_range
+        scores = list(np.searchsorted(sentence_starts, [e.start + 1 for e in entities]))
         scores = [sentences[i - 1][1] for i in scores]
-        scores = [scores[i] + int(10 * entities[i][2]) for i in range(len(entities))]
+        scores = [scores[i] + 10 * entities[i].score for i in range(len(entities))]
         for i in range(len(entities)):
-            entities[i] = (entities[i][0], entities[i][1], scores[i], entities[i][3])
+            entities[i].score = scores[i]
+    else:
+        for i in range(len(entities)):
+            entities[i].score = 0.0
+    if with_comentions:
         for i in range(len(entities)):
             entity = entities[i]
-            count = 1
             comentions = [
-                entities[j][0]
+                entities[j].name
                 for j in range(len(entities))
-                if j != i and abs(entities[j][3] - entity[3]) < math.ceil(context_width / 2)
+                if j != i and abs(entities[j].start - entity.start) < math.ceil(context_width / 2)
             ]
-            entities[i] = (
-                entity[0],
-                entity[1],
-                entity[2],
-                entity[3],
-                count,
-                comentions,
-            )
+            entities[i].comentions = comentions
+    if with_context:
         for i in range(len(entities)):
             entity = entities[i]
-            if entity[3] >= 0 and entity[3] < len(text):
-                left = max(0, entity[3] - math.floor(context_width / 2))
-                right = min(len(text), entity[3] + math.ceil(context_width / 2))
+            if entity.start >= 0 and entity.start < len(text):
+                left = max(0, entity.start - math.floor(context_width / 2))
+                right = min(len(text), entity.start + math.ceil(context_width / 2))
                 context = ("[..]" if left > 0 else "") + text[left:right] + ("[..]" if right < len(text) else "")
-                entities[i] = (
-                    entity[0],
-                    entity[1],
-                    entity[2],
-                    context,
-                    entity[4],
-                    entity[5],
-                )
-        entities = [NERObject(*entities[i]) for i in range(len(entities))]
+                entities[i].context = context
     return entities
 
 
@@ -264,33 +299,36 @@ def get_extractive_summary(text, language, max_chars, fast=False, with_scores=Fa
     return summary
 
 
-def ner_pipe(text, language, spacy_model, flair_model=None, fast=False, compression_ratio="auto"):
+def ner_pipe(
+    text,
+    language,
+    spacy_model,
+    flair_model=None,
+    fast=False,
+    compression_ratio="auto",
+    with_scores=True,
+    with_comentions=True,
+    with_context=True,
+):
     if compression_ratio == "auto":
         compression_ratio = max(1.0, len(text) / 15000) if fast else 1.0
     sentences = get_extractive_summary(text, language, int(len(text) / compression_ratio), fast=fast, with_scores=True)
-    ner = compute_ner(language, sentences, spacy_model, flair_model)
+    ner = compute_ner(language, sentences, spacy_model, flair_model, 150, with_scores, with_comentions, with_context)
     return ner
 
 
-def get_ner_handler(language, fast=False, compression_ratio="auto"):
+def get_ner_handler(language, fast=False):
     try:
         get_nltk_tokenizer(language)  # raises a LookupError if the language is not valid
     except LookupError:
         language = "english"
-    spacy_model = SPACY_NER_MODELS.get(language, SPACY_NER_MODELS['english'])()
-    flair_model = None if fast else FLAIR_NER_MODELS.get(language, FLAIR_NER_MODELS['english'])()
-    return lambda text: ner_pipe(text, language, spacy_model, flair_model, fast, compression_ratio)
+    spacy_model = SPACY_NER_MODELS.get(language, SPACY_NER_MODELS["english"])()
+    flair_model = None if fast else FLAIR_NER_MODELS.get(language, FLAIR_NER_MODELS["english"])()
+    return lambda text, compression_ratio="auto", with_scores=True, with_comentions=True, with_context=True: ner_pipe(
+        text, language, spacy_model, flair_model, fast, compression_ratio, with_scores, with_comentions, with_context
+    )
 
 
 @st.cache_resource
 def get_cached_ner_handler(language, fast):
     return get_ner_handler(language, fast)
-
-
-def test():
-    text = """My name is Valerio Simoni, and I live in NYC. I love the Colosseum, and my phone is +123 456 789.
-my email is aaa@ggg.com, but my ip address is secret! 123.123.123.0:1111"""
-    entities = get_ner_handler("english", True)(text)
-    print(entities)
-    
-test()