streamlit-octostar-utils 0.1.7a3__tar.gz → 0.1.7a5__tar.gz

{streamlit_octostar_utils-0.1.7a3 → streamlit_octostar_utils-0.1.7a5}/PKG-INFO

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

{streamlit_octostar_utils-0.1.7a3 → streamlit_octostar_utils-0.1.7a5}/pyproject.toml

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

{streamlit_octostar_utils-0.1.7a3 → streamlit_octostar_utils-0.1.7a5}/streamlit_octostar_utils/api_crafter/celery.py

@@ -23,7 +23,7 @@ from functools import wraps
 import logging
 
 logger = logging.getLogger(__name__)
-logging.getLogger('pottery').setLevel(logging.WARNING)
+logging.getLogger("pottery").setLevel(logging.WARNING)
 from celery.app.defaults import DEFAULTS as CELERY_DEFAULTS
 import urllib
 
@@ -630,6 +630,11 @@ class FastAPICeleryTaskRoute(Route):
                 except BaseException as e:
                     exc = e
             data = {}
+            assert (
+                (state in ["FAILURE", "RETRY", "REVOKED"] and exc is not None)
+                or (state in ["SUCCESS"] and exc is None)
+                or (state not in ["SUCCESS", "FAILURE", "RETRY", "REVOKED"])
+            )
             if state in ["FAILURE", "RETRY", "REVOKED"]:
                 error_response = DefaultErrorRoute.format_error(exc, debug=True).body.decode("utf-8")
                 data = {
@@ -682,11 +687,16 @@ class CeleryErrorRoute(DefaultErrorRoute):
         debug=False,
         excs_to_status_codes=None,
         silenced_excs=None,
-        log_filter=None
-        ,
+        log_filter=None,
     ):
         if excs_to_status_codes is None:
-            excs_to_status_codes = {**DefaultErrorRoute.DEFAULT_STATUS_CODE_MAPPINGS, **CeleryErrorRoute.DEFAULT_STATUS_CODE_MAPPINGS}
+            excs_to_status_codes = {
+                **DefaultErrorRoute.DEFAULT_STATUS_CODE_MAPPINGS,
+                **CeleryErrorRoute.DEFAULT_STATUS_CODE_MAPPINGS,
+            }
         if silenced_excs is None:
-            silenced_excs = {**DefaultErrorRoute.DEFAULT_SILENCED_EXCEPTIONS, **CeleryErrorRoute.DEFAULT_SILENCED_EXCEPTIONS}
+            silenced_excs = {
+                **DefaultErrorRoute.DEFAULT_SILENCED_EXCEPTIONS,
+                **CeleryErrorRoute.DEFAULT_SILENCED_EXCEPTIONS,
+            }
         DefaultErrorRoute.add_default_exceptions_handler(fs_app, debug, excs_to_status_codes)

{streamlit_octostar_utils-0.1.7a3 → streamlit_octostar_utils-0.1.7a5}/streamlit_octostar_utils/api_crafter/nifi.py

@@ -1103,14 +1103,20 @@ class NifiRoute(Route):
     def define_routes(self):
         @Route.route(self, path="/task-state/{task_id}")
         async def get_task_status(task_id: str) -> JSONResponse:
-            task_status = await self.tasks_routes.get_task(task_id, pop=False)
-            task_status = task_status.model_dump(mode="json")["data"]["task_state"]
+            try:
+                task_status = await self.tasks_routes.get_task(task_id, pop=False)
+                task_status = task_status.model_dump(mode="json")["data"]["task_state"]
+            except BaseException as e:
+                raise ValueError(f"Could not fetch task state for task id {task_id}!\n{e}")
             return JSONResponse(task_status)
 
         @Route.route(self, path="/task-result/{task_id}")
         async def get_task_result(task_id: str) -> JSONResponse:
-            return_data = await self.tasks_routes.get_task(task_id, pop=True)
-            return_data = return_data.model_dump(mode="json")["data"]["data"]
+            try:
+                return_data = await self.tasks_routes.get_task(task_id, pop=True)
+                return_data = return_data.model_dump(mode="json")["data"]["data"]
+            except BaseException as e:
+                raise ValueError(f"Could not fetch task result for task id {task_id}\n{e}!")
             return JSONResponse(return_data)
 
         @Route.route(self, path="/{op}", methods=["POST"])

{streamlit_octostar_utils-0.1.7a3 → streamlit_octostar_utils-0.1.7a5}/streamlit_octostar_utils/nlp/language.py

@@ -12,4 +12,4 @@ def detect_language(text, min_confidence=None):
         return None
     detected_lang = re.sub("[^A-Za-z]", "", detected_lang).lower()
     detected_lang = languages.to_name(detected_lang).lower()
-    return detected_lang
+    return detected_lang, confidence

{streamlit_octostar_utils-0.1.7a3 → streamlit_octostar_utils-0.1.7a5}/streamlit_octostar_utils/nlp/ner.py

@@ -1,7 +1,5 @@
 import re
 import streamlit as st
-import py3langid as langid
-import iso639 as languages
 from spacy_download import load_spacy
 from flair.data import Sentence
 from flair.models import SequenceTagger
@@ -17,8 +15,6 @@ import math
 import nltk
 from typing import Optional, List
 
-nltk.download("punkt")
-
 SPACY_NER_MODELS = {
     "english": lambda: load_spacy(
         "en_core_web_sm",
@@ -26,9 +22,17 @@ 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))"
+    ],
+    "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))",
+}
 
 BASE_TO_ONTONOTES_LABELMAP = {"PER": "PERSON"}
-BASE_ALLOWED_LABELS = ["PERSON", "ORG", "LOC", "NORP", "GPE", "PRODUCT", "DATE", "TIME"]
+BASE_ALLOWED_LABELS = ["PERSON", "ORG", "LOC", "NORP", "GPE", "PRODUCT", "DATE", "TIME", "PHONE", "IP_ADDRESS", "EMAIL"]
 
 
 def _sumy__get_best_sentences(sentences, rating, *args, **kwargs):
@@ -39,9 +43,7 @@ def _sumy__get_best_sentences(sentences, rating, *args, **kwargs):
     if isinstance(rating, dict):
         assert not args and not kwargs
         rate = lambda s: rating[s]
-    infos = (
-        SentenceInfo(s, o, rate(s, *args, **kwargs)) for o, s in enumerate(sentences)
-    )
+    infos = (SentenceInfo(s, o, rate(s, *args, **kwargs)) for o, s in enumerate(sentences))
     infos = sorted(infos, key=attrgetter("rating"), reverse=True)
     return tuple((i.sentence, i.rating, i.order) for i in infos)
 
@@ -62,9 +64,15 @@ def _sumy__lsa_call(summarizer, document):
 
 def _sumy__luhn_call(summarizer, document):
     words = summarizer._get_significant_words(document.words)
-    return _sumy__get_best_sentences(
-        document.sentences, summarizer.rate_sentence, words
-    )
+    return _sumy__get_best_sentences(document.sentences, summarizer.rate_sentence, words)
+
+
+def get_nltk_tokenizer(language: str) -> Tokenizer:
+    try:
+        nltk.data.find("tokenizers/punkt")
+    except LookupError:
+        nltk.download("punkt")
+    return Tokenizer(language)
 
 
 class NERObject(object):
@@ -77,22 +85,26 @@ class NERObject(object):
         self.comentions: Optional[List[str]] = comentions
         self.sources: Optional[List[str]] = list()
 
-
-def detect_language(text, min_confidence=None):
-    detector = langid.langid.LanguageIdentifier.from_pickled_model(
-        langid.langid.MODEL_FILE, norm_probs=True
-    )
-    detected_lang, confidence = detector.classify(text)
-    if min_confidence and confidence < min_confidence:
-        return None
-    detected_lang = re.sub("[^A-Za-z]", "", detected_lang).lower()
-    detected_lang = languages.to_name(detected_lang).lower()
-    return detected_lang
+    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
+        return data
+            
+    def __repr__(self):
+        return f"NERObject(label={self.label},name={self.name})"
 
 
-def postprocess_ner(entities, allowed_labels, max_entities=100):
-    if allowed_labels != "all":
-        entities = [e for e in entities if e.label in allowed_labels]
+def postprocess_ner(entities, 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)
     final_entities = []
     for _, group in itertools.groupby(entities, key=lambda x: x.name):
@@ -108,34 +120,24 @@ def postprocess_ner(entities, allowed_labels, max_entities=100):
         )
         best_entity.sources = list(set(itertools.chain(*[e.sources for e in group])))
         final_entities.append(best_entity)
-    final_entities = sorted(
-        final_entities, key=lambda x: x.score * x.count, reverse=True
-    )
-    if len(final_entities) > max_entities:
+    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, fast=True, context_width=150):
+def compute_ner(language, sentences, spacy_model, flair_model=None, context_width=150):
     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])
-    if fast:
-        model = SPACY_NER_MODELS.get(language, SPACY_NER_MODELS["english"])()
-        entities = [
-            (
-                entity.text,
-                BASE_TO_ONTONOTES_LABELMAP.get(entity.label_, entity.label_),
-                0,
-                entity.start_char,
-            )
-            for entity in model(text).ents
-        ]
-    else:
-        model = FLAIR_NER_MODELS.get(language, FLAIR_NER_MODELS["english"])()
+    min_score = 1.0
+    entities = []
+
+    # FLAIR model (if not fast)
+    if flair_model:
         input = [Sentence(sentence[0]) for sentence in sentences]
-        model.predict(input)
+        flair_model.predict(input)
         output = [e for sentence in input for e in sentence.get_spans("ner")]
         flair_entities = [
             (
@@ -145,35 +147,45 @@ def compute_ner(language, sentences, fast=True, context_width=150):
                     entity.annotation_layers["ner"][0].value,
                 ),
                 entity.score,
-                sentence_starts[input.index(entity[0].sentence)]
-                + entity[0].start_position,
+                sentence_starts[input.index(entity[0].sentence)] + entity[0].start_position,
             )
             for entity in output
         ]
-        min_score = min([0] + [e[2] for e in flair_entities])
-        model = SPACY_NER_MODELS.get(language, SPACY_NER_MODELS["english"])()
-        spacy_entities = [
-            (
-                entity.text,
-                BASE_TO_ONTONOTES_LABELMAP.get(entity.label_, entity.label_),
-                min_score - 1,
-                entity.start_char,
-            )
-            for entity in model(text).ents
-        ]
-        entities = flair_entities + spacy_entities
+        min_score = min(min_score, *[e[2] for e in flair_entities])
+        entities += flair_entities
+        del flair_entities
+
+    # 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)
+            ]
+            print(regex_entities)
+            entities += 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
+
+    # 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])
-        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
-        ]
+        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]))
         scores = [sentences[i - 1][1] for i in scores]
         scores = [scores[i] + int(10 * entities[i][2]) for i in range(len(entities))]
@@ -185,8 +197,7 @@ def compute_ner(language, sentences, fast=True, context_width=150):
             comentions = [
                 entities[j][0]
                 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][3] - entity[3]) < math.ceil(context_width / 2)
             ]
             entities[i] = (
                 entity[0],
@@ -201,11 +212,7 @@ def compute_ner(language, sentences, fast=True, context_width=150):
             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))
-                context = (
-                    ("[..]" if left > 0 else "")
-                    + text[left:right]
-                    + ("[..]" if right < len(text) else "")
-                )
+                context = ("[..]" if left > 0 else "") + text[left:right] + ("[..]" if right < len(text) else "")
                 entities[i] = (
                     entity[0],
                     entity[1],
@@ -214,22 +221,12 @@ def compute_ner(language, sentences, fast=True, context_width=150):
                     entity[4],
                     entity[5],
                 )
-        entities = [
-            NERObject(
-                entities[i][0],
-                entities[i][1],
-                entities[i][2],
-                entities[i][3],
-                entities[i][4],
-                entities[i][5],
-            )
-            for i in range(len(entities))
-        ]
+        entities = [NERObject(*entities[i]) for i in range(len(entities))]
     return entities
 
 
 def get_extractive_summary(text, language, max_chars, fast=False, with_scores=False):
-    tokenizer = Tokenizer(language)
+    tokenizer = get_nltk_tokenizer(language)
     stemmer = Stemmer(language)
     parser = PlaintextParser.from_string(text, tokenizer)
     if fast:
@@ -268,24 +265,24 @@ def get_extractive_summary(text, language, max_chars, fast=False, with_scores=Fa
     return summary
 
 
-def ner_pipe(text, language, fast=False, compression_ratio="auto"):
+def ner_pipe(text, language, spacy_model, flair_model=None, fast=False, compression_ratio="auto"):
     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, fast=fast)
+    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)
     return ner
 
 
 def get_ner_handler(language, fast=False, compression_ratio="auto"):
     try:
-        Tokenizer(language)  # raises a LookupError if the language is not valid
+        get_nltk_tokenizer(language)  # raises a LookupError if the language is not valid
     except LookupError:
         language = "english"
-    return lambda text: ner_pipe(text, language, 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: ner_pipe(text, language, spacy_model, flair_model, fast, compression_ratio)
 
 
 @st.cache_resource
 def get_cached_ner_handler(language, fast):
-    return get_ner_handler(language, fast)
+    return get_ner_handler(language, fast)