streamlit-octostar-utils 0.2.10__tar.gz → 2.11a2__tar.gz

{streamlit_octostar_utils-0.2.10 → streamlit_octostar_utils-2.11a2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: streamlit-octostar-utils
-Version: 0.2.10
+Version: 2.11a2
 Summary: 
 License: MIT
 License-File: LICENSE

{streamlit_octostar_utils-0.2.10 → streamlit_octostar_utils-2.11a2}/pyproject.toml

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

{streamlit_octostar_utils-0.2.10 → streamlit_octostar_utils-2.11a2}/streamlit_octostar_utils/api_crafter/fastapi.py

@@ -330,7 +330,7 @@ class DefaultErrorRoute:
         if len(message) > MAX_ERROR_MESSAGE_BYTES:
             message = message[-MAX_ERROR_MESSAGE_BYTES:]
         try:
-            tcbk = "\n".join(traceback.format_exception(exc))
+            tcbk = traceback.format_exception(exc)
             if len(tcbk) > MAX_ERROR_TRACEBACK_BYTES:
                 tcbk = tcbk[-MAX_ERROR_TRACEBACK_BYTES:]
         except:

streamlit_octostar_utils-2.11a2/streamlit_octostar_utils/nlp/language.py

@@ -0,0 +1,55 @@
+import re
+import py3langid as langid
+import iso639 as languages
+
+
+def alpha2_to_language(alpha2: str) -> str:
+    if not alpha2:
+        return None
+    code = alpha2.strip().lower()
+    return languages.to_name(code)
+
+def language_to_alpha2(language_name: str) -> str:
+    if not language_name:
+        return None
+    name = language_name.strip().lower()
+    data = languages.find(name)
+    return data["iso639_1"]
+
+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, confidence
+    detected_lang = re.sub("[^A-Za-z]", "", detected_lang).lower()
+    detected_lang = languages.to_name(detected_lang).lower()
+    return detected_lang, confidence
+
+FLAIR_MODELS = {
+    "en": "flair/ner-english-large",
+    "es": "flair/ner-spanish-large",
+    "de": "flair/ner-german-large",
+    "nl": "flair/ner-dutch-large",
+    "multi": "flair/ner-multi",
+    "multi-fast": "flair/ner-multi-fast",
+}
+
+SPACY_MODELS = {
+    "en": 'en_core_web_sm',
+}
+
+def load_language_model(language, type):
+    from flair.models import SequenceTagger
+    from spacy_download import load_spacy
+
+    model = None
+    match type:
+        case "spacy":
+            model_name = SPACY_MODELS.get(language, SPACY_MODELS["en"])
+            model = load_spacy(model_name)
+        case "flair":
+            model_name = FLAIR_MODELS.get(language, "flair/ner-multi")
+            model = SequenceTagger.load(model_name)
+    return model

streamlit_octostar_utils-2.11a2/streamlit_octostar_utils/nlp/ner.py

@@ -0,0 +1,634 @@
+import itertools
+import math
+from typing import Optional, List, Tuple
+from pydantic import BaseModel, ConfigDict, Field
+from collections import Counter
+
+from presidio_analyzer import AnalyzerEngine, BatchAnalyzerEngine, RecognizerRegistry, AnalysisExplanation, \
+    EntityRecognizer, RecognizerResult
+from presidio_analyzer.nlp_engine import NlpArtifacts, NlpEngineProvider
+import streamlit as st
+import nltk
+import pandas as pd
+from flair.data import Sentence
+from flair.models import SequenceTagger
+
+from sumy.parsers.plaintext import PlaintextParser
+from sumy.nlp.tokenizers import Tokenizer
+from sumy.nlp.stemmers import Stemmer
+from sumy.summarizers.lsa import LsaSummarizer
+from sumy.summarizers.luhn import LuhnSummarizer
+from sumy.utils import get_stop_words
+
+from nlp.language import alpha2_to_language
+
+BASE_ALLOWED_LABELS = ["PERSON", "ORG", "LOC", "NORP", "GPE", "PRODUCT", "DATE", "PHONE", "IP_ADDRESS", "EMAIL", "URL",
+                       "CRYPTO", "IBAN", "CREDIT_CARD", "US_SSN", "US_DRIVER_LICENSE", "US_PASSPORT", "MEDICAL_LICENSE"]
+
+PRESIDIO_TO_BASE_ALIASES = {
+    "PHONE_NUMBER": "PHONE",
+    "EMAIL_ADDRESS": "EMAIL",
+    "IBAN_CODE": "IBAN",
+    "DRIVER_LICENSE": "US_DRIVER_LICENSE",
+    "US_DRIVER_LICENSE": "US_DRIVER_LICENSE",
+    "US_DRIVERS_LICENSE": "US_DRIVER_LICENSE",
+    "PASSPORT": "US_PASSPORT",
+    "CREDIT_CARD": "CREDIT_CARD",
+    "URL": "URL",
+    "IP_ADDRESS": "IP_ADDRESS",
+    "CRYPTO": "CRYPTO",
+    "CRYPTO_WALLET": "CRYPTO",
+    "CRYPTO_WALLET_ADDRESS": "CRYPTO",
+    "DATE_TIME": "DATE",
+    "LOCATION": "LOC",
+    "ORGANIZATION": "ORG",
+}
+
+BASE_TO_RECOGNIZER_EXPANSIONS = {
+    "ORG": ["ORG", "ORGANIZATION"],
+    "LOC": ["LOC", "LOCATION"],
+    "PHONE": ["PHONE", "PHONE_NUMBER"],
+    "EMAIL": ["EMAIL", "EMAIL_ADDRESS"],
+    "IBAN": ["IBAN", "IBAN_CODE"],
+    "US_DRIVER_LICENSE": ["US_DRIVER_LICENSE", "US_DRIVERS_LICENSE", "DRIVER_LICENSE"],
+    "US_PASSPORT": ["US_PASSPORT", "PASSPORT"],
+    "DATE": ["DATE", "DATE_TIME"],
+    "PERSON": ["PERSON"],
+    "URL": ["URL"],
+    "IP_ADDRESS": ["IP_ADDRESS"],
+    "CRYPTO": ["CRYPTO", "CRYPTO_WALLET", "CRYPTO_WALLET_ADDRESS"],
+    "CREDIT_CARD": ["CREDIT_CARD"],
+    "US_SSN": ["US_SSN"],
+    "MEDICAL_LICENSE": ["MEDICAL_LICENSE"],
+    "NORP": ["NORP"],
+    "GPE": ["GPE"],
+    "PRODUCT": ["PRODUCT"],
+}
+
+BASE_TO_ONTONOTES_LABELMAP = {"PER": "PERSON"}
+
+
+class FlairRecognizer(EntityRecognizer):
+    ENTITIES = [
+        "LOC",
+        "PERSON",
+        "ORG",
+    ]
+
+    DEFAULT_EXPLANATION = "Identified as {} by Flair's Named Entity Recognition"
+
+    CHECK_LABEL_GROUPS = [
+        ({"LOC"}, {"LOC", "LOCATION"}),
+        ({"PERSON"}, {"PER", "PERSON"}),
+        ({"ORG"}, {"ORG", "ORGANIZATION"}),
+    ]
+
+    MODEL_LANGUAGES = {
+        "en": "flair/ner-english-large",
+        "es": "flair/ner-spanish-large",
+        "de": "flair/ner-german-large",
+        "nl": "flair/ner-dutch-large",
+        "multi": "flair/ner-multi",
+        "multi-fast": "flair/ner-multi-fast",
+    }
+
+    PRESIDIO_EQUIVALENCES = {
+        "PER": "PERSON",
+        "LOC": "LOC",
+        "ORG": "ORG"
+    }
+
+    def __init__(
+            self,
+            model: SequenceTagger = None,
+            supported_language: str = "en",
+            supported_entities: Optional[List[str]] = None,
+            check_label_groups: Optional[Tuple[set, set]] = None,
+    ):
+        self.check_label_groups = (
+            check_label_groups if check_label_groups else self.CHECK_LABEL_GROUPS
+        )
+
+        supported_entities = supported_entities if supported_entities else self.ENTITIES
+        self.model = model
+
+        super().__init__(
+            supported_entities=supported_entities,
+            supported_language=supported_language,
+            name="Flair Analytics",
+        )
+
+    def load(self) -> None:
+        pass
+
+    def get_supported_entities(self) -> List[str]:
+        return self.supported_entities
+
+    def analyze(self, text: str, entities: List[str], nlp_artifacts: NlpArtifacts = None) -> List[RecognizerResult]:
+        results = []
+
+        sentences = Sentence(text)
+        self.model.predict(sentences)
+
+        if not entities:
+            entities = self.supported_entities
+
+        for entity in entities:
+            if entity not in self.supported_entities:
+                continue
+
+            for ent in sentences.get_spans("ner"):
+                if not self.__check_label(
+                        entity, ent.labels[0].value, self.check_label_groups
+                ):
+                    continue
+                textual_explanation = self.DEFAULT_EXPLANATION.format(
+                    ent.labels[0].value
+                )
+                explanation = self.build_flair_explanation(
+                    round(ent.score, 2), textual_explanation
+                )
+                flair_result = self._convert_to_recognizer_result(ent, explanation)
+
+                results.append(flair_result)
+
+        return results
+
+    def build_flair_explanation(self, original_score: float, explanation: str) -> AnalysisExplanation:
+        explanation = AnalysisExplanation(
+            recognizer=self.__class__.__name__,
+            original_score=original_score,
+            textual_explanation=explanation,
+        )
+        return explanation
+
+    def _convert_to_recognizer_result(self, entity, explanation) -> RecognizerResult:
+        entity_type = self.PRESIDIO_EQUIVALENCES.get(entity.tag, entity.tag)
+        flair_score = round(entity.score, 2)
+
+        flair_results = RecognizerResult(
+            entity_type=entity_type,
+            start=entity.start_position,
+            end=entity.end_position,
+            score=flair_score,
+            analysis_explanation=explanation,
+        )
+
+        return flair_results
+
+    @staticmethod
+    def __check_label(
+            entity: str, label: str, check_label_groups: Tuple[set, set]
+    ) -> bool:
+        return any(
+            [entity in egrp and label in lgrp for egrp, lgrp in check_label_groups]
+        )
+
+
+def normalize_label(label: str) -> str:
+    return PRESIDIO_TO_BASE_ALIASES.get(label, label)
+
+
+def expand_entities_for_analyzer(entities_list):
+    expanded = set()
+    for e in entities_list:
+        vals = BASE_TO_RECOGNIZER_EXPANSIONS.get(e, [e])
+        for v in vals:
+            expanded.add(v)
+    return list(expanded)
+
+
+def _sumy__get_best_sentences(sentences, rating, *args, **kwargs):
+    from operator import attrgetter
+    from sumy.summarizers._summarizer import SentenceInfo
+
+    rate = rating
+    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 = sorted(infos, key=attrgetter("rating"), reverse=True)
+    return tuple((i.sentence, i.rating, i.order) for i in infos)
+
+
+def _sumy__lsa_call(summarizer, document):
+    summarizer._ensure_dependecies_installed()
+    dictionary = summarizer._create_dictionary(document)
+    if not dictionary:
+        return ()
+    matrix = summarizer._create_matrix(document, dictionary)
+    matrix = summarizer._compute_term_frequency(matrix)
+    from numpy.linalg import svd as singular_value_decomposition
+
+    u, sigma, v = singular_value_decomposition(matrix, full_matrices=False)
+    ranks = iter(summarizer._compute_ranks(sigma, v))
+    return _sumy__get_best_sentences(document.sentences, lambda s: next(ranks))
+
+
+def _sumy__luhn_call(summarizer, document):
+    words = summarizer._get_significant_words(document.words)
+    return _sumy__get_best_sentences(document.sentences, summarizer.rate_sentence, words)
+
+
+def get_nltk_tokenizer(language: str) -> Tokenizer:
+    nltk_lang = alpha2_to_language(language).lower()
+
+    try:
+        nltk.data.find("tokenizers/punkt")
+    except LookupError:
+        nltk.download("punkt")
+
+    return Tokenizer(nltk_lang)
+
+
+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 __repr__(self):
+        return f"NERObject(label={self.label},name={self.name})"
+
+
+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)
+    final_entities = []
+    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)
+        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 build_presidio_analyzer(language: str, engine_type: str = "spacy", model=None) -> AnalyzerEngine:
+    registry = RecognizerRegistry()
+
+    if engine_type == "flair":
+
+        flair_recognizer = FlairRecognizer(
+            model=model,
+            supported_language=language
+        )
+        registry.add_recognizer(flair_recognizer)
+
+        default_registry = RecognizerRegistry()
+        default_registry.load_predefined_recognizers()
+
+        flair_handled_entities = {"PERSON", "LOC", "ORG"}
+
+        for recognizer in default_registry.recognizers:
+            recognizer_entities = set(recognizer.supported_entities) if hasattr(recognizer, 'supported_entities') else set()
+
+            if recognizer_entities and recognizer_entities.issubset(flair_handled_entities):
+                continue
+
+            registry.add_recognizer(recognizer)
+
+        return AnalyzerEngine(
+            registry=registry,
+            supported_languages=[language]
+        )
+
+    else:
+        registry.load_predefined_recognizers()
+
+        if model is None:
+            raise ValueError("SpaCy model name must be provided")
+
+        configuration = {
+            "nlp_engine_name": "spacy",
+            "models": [{"lang_code": language, "model_name": model}],
+        }
+
+        provider = NlpEngineProvider(nlp_configuration=configuration)
+        nlp_engine = provider.create_engine()
+
+        return AnalyzerEngine(
+            nlp_engine=nlp_engine,
+            registry=registry,
+            supported_languages=[language],
+        )
+
+
+def analyze_column_sample(column_values: pd.Series, analyzer: AnalyzerEngine, language: str,
+                          entities: Optional[List[str]], score_threshold: float) -> Optional[str]:
+    sample_values = column_values.dropna().head(50)
+
+    if sample_values.empty:
+        return None
+
+    entity_counter = Counter()
+
+    for value in sample_values:
+        text = str(value).strip()
+
+        if not text:
+            continue
+
+        results = analyzer.analyze(
+            text=text,
+            language=language,
+            entities=(expand_entities_for_analyzer(entities) if entities else None)
+        )
+
+        for result in results:
+            if result.score >= score_threshold:
+                entity_counter[normalize_label(result.entity_type)] += 1
+
+    if not entity_counter:
+        return None
+
+    most_common = entity_counter.most_common(1)[0]
+    total_detections = sum(entity_counter.values())
+
+    if most_common[1] > total_detections * 0.5:
+        return most_common[0]
+
+    return most_common[0] if entity_counter else None
+
+
+def analyze_dataframe_optimized(df: pd.DataFrame, analyzer: AnalyzerEngine, language: str,
+                                entities: Optional[List[str]] = None, score_threshold: float = 0.5) -> List[NERObject]:
+    ner_objects = []
+
+    for column_name in df.columns:
+        entity_type = analyze_column_sample(
+            df[column_name],
+            analyzer,
+            language,
+            entities,
+            score_threshold
+        )
+
+        if entity_type:
+            for idx, value in df[column_name].dropna().items():
+                text = str(value).strip()
+
+                if text:
+                    ner_objects.append(NERObject(
+                        name=text[:100],
+                        label=entity_type,
+                        score=0.9,
+                        start=0,
+                        count=1,
+                        context=text[:100]
+                    ))
+
+    return ner_objects
+
+
+def compute_ner_presidio(
+        text,
+        language,
+        analyzer,
+        entities=None,
+        score_threshold=0.5,
+        context_width=150,
+        with_comentions=True,
+        with_context=True,
+        batch_size=32,
+        n_process=4
+):
+    if isinstance(text, pd.DataFrame):
+        if len(text) >= 100:
+            return analyze_dataframe_optimized(text, analyzer, language, entities, score_threshold)
+
+        else:
+            texts = []
+
+            for col in text.columns:
+                for idx, value in text[col].dropna().items():
+                    text_value = str(value).strip()
+
+                    if text_value:
+                        texts.append(text_value)
+
+            text = "\n".join(texts)
+
+    elif isinstance(text, list):
+        batch_analyzer = BatchAnalyzerEngine(analyzer_engine=analyzer)
+
+        results_generator = batch_analyzer.analyze_iterator(
+            texts=text,
+            language=language,
+            batch_size=batch_size,
+            n_process=n_process,
+            entities=(expand_entities_for_analyzer(entities) if entities else None),
+        )
+
+        all_results = list(results_generator)
+        ner_objects = []
+
+        for text_item, results in zip(text, all_results):
+            for result in results:
+                if result.score >= score_threshold:
+                    context_start = max(0, result.start - 30)
+                    context_end = min(len(text_item), result.end + 30)
+                    context = text_item[context_start:context_end] if with_context else None
+
+                    ner_objects.append(NERObject(
+                        name=text_item[result.start:result.end],
+                        label=normalize_label(result.entity_type),
+                        score=float(result.score),
+                        start=int(result.start),
+                        count=1,
+                        context=context
+                    ))
+
+        return ner_objects
+
+    results = analyzer.analyze(
+        text=text,
+        language=language,
+        entities=(expand_entities_for_analyzer(entities) if entities else None)
+    )
+
+    ner_objects = []
+
+    for result in results:
+        if result.score >= score_threshold:
+            context_start = max(0, result.start - math.floor(context_width / 2))
+            context_end = min(len(text), result.end + math.ceil(context_width / 2))
+            context = text[context_start:context_end] if with_context else None
+
+            ner_objects.append(NERObject(
+                name=text[result.start:result.end],
+                label=normalize_label(result.entity_type),
+                score=float(result.score),
+                start=int(result.start),
+                count=1,
+                context=context
+            ))
+
+    if with_comentions:
+        for i in range(len(ner_objects)):
+            entity = ner_objects[i]
+            comentions = [
+                ner_objects[j].name
+                for j in range(len(ner_objects))
+                if j != i and abs(ner_objects[j].start - entity.start) < math.ceil(context_width / 2)
+            ]
+            ner_objects[i].comentions = comentions
+
+    return ner_objects
+
+
+def get_extractive_summary(text, language, max_chars, fast=False, with_scores=False):
+    tokenizer = get_nltk_tokenizer(language)
+    stemmer = Stemmer(language)
+    parser = PlaintextParser.from_string(text, tokenizer)
+    if fast:
+        summarizer = LuhnSummarizer(stemmer)
+        summarizer.stop_words = get_stop_words(language)
+        scored_sentences = iter(_sumy__luhn_call(summarizer, parser.document))
+    else:
+        summarizer = LsaSummarizer(stemmer)
+        summarizer.stop_words = get_stop_words(language)
+        scored_sentences = iter(_sumy__lsa_call(summarizer, parser.document))
+    summary = []
+    summary_chars = 0
+    summary_chars_penultimate = 0
+    while summary_chars < max_chars:
+        try:
+            next_sentence = next(scored_sentences)
+            summary.append(next_sentence)
+            summary_chars_penultimate = summary_chars
+            summary_chars += len(" " + next_sentence[0]._text)
+        except StopIteration:
+            break
+    summary = sorted(summary, key=lambda x: x[2])
+    summary = [(sentence[0]._text, sentence[1]) for sentence in summary]
+    if summary_chars > max_chars:
+        summary[-1] = (
+            summary[-1][0][: max_chars - summary_chars_penultimate],
+            summary[-1][1],
+        )
+    if not with_scores:
+        summary = " ".join([s[0] for s in summary])
+    else:
+        min_score = min([s[1] for s in summary]) if summary else 0
+        max_score = max([min_score] + [s[1] for s in summary])
+        score_range = 1 if min_score == max_score else (max_score - min_score)
+        summary = [(s[0], (s[1] - min_score) / score_range) for s in summary]
+    return summary
+
+
+def ner_pipe(
+        text,
+        language,
+        model,
+        engine_type="spacy",
+        fast=False,
+        compression_ratio="auto",
+        with_comentions=True,
+        with_context=True,
+        entities=None,
+        score_threshold=0.5,
+        batch_size=32,
+        n_process=4
+):
+    analyzer = build_presidio_analyzer(
+        language=language,
+        engine_type=engine_type,
+        model=model,
+    )
+
+    if isinstance(text, pd.DataFrame):
+        ner = compute_ner_presidio(
+            text,
+            language,
+            analyzer,
+            entities,
+            score_threshold,
+            with_comentions=with_comentions,
+            with_context=with_context,
+            batch_size=batch_size,
+            n_process=n_process
+        )
+    else:
+        if compression_ratio == "auto":
+            compression_ratio = max(1.0, len(text) / 15000) if fast else 1.0
+
+        if compression_ratio > 1.0:
+            sentences = get_extractive_summary(text, language, int(len(text) / compression_ratio), fast=fast,
+                                               with_scores=True)
+            text = " ".join([s[0] for s in sentences])
+
+        ner = compute_ner_presidio(
+            text,
+            language,
+            analyzer,
+            entities,
+            score_threshold,
+            with_comentions=with_comentions,
+            with_context=with_context,
+            batch_size=batch_size,
+            n_process=n_process
+        )
+
+    return ner
+
+
+def get_ner_handler(
+        language,
+        model,
+        engine_type="spacy",
+        fast=False,
+        entities=None,
+        score_threshold=0.5,
+        batch_size=32,
+        n_process=4
+):
+    try:
+        get_nltk_tokenizer(language)
+    except LookupError:
+        language = "en"
+
+    return lambda text, compression_ratio="auto", with_comentions=True, with_context=True: ner_pipe(
+        text,
+        language,
+        model,
+        engine_type,
+        fast,
+        compression_ratio,
+        with_comentions,
+        with_context,
+        entities,
+        score_threshold,
+        batch_size,
+        n_process
+    )
+
+
+@st.cache_resource
+def get_cached_ner_handler(language, model):
+    return get_ner_handler(language, model)

streamlit_octostar_utils-0.2.10/streamlit_octostar_utils/nlp/language.py

@@ -1,15 +0,0 @@
-import re
-import py3langid as langid
-import iso639 as languages
-
-
-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, confidence
-    detected_lang = re.sub("[^A-Za-z]", "", detected_lang).lower()
-    detected_lang = languages.to_name(detected_lang).lower()
-    return detected_lang, confidence

streamlit_octostar_utils-0.2.10/streamlit_octostar_utils/nlp/ner.py

@@ -1,329 +0,0 @@
-import re
-import streamlit as st
-from spacy_download import load_spacy
-from flair.data import Sentence
-from flair.models import SequenceTagger
-from sumy.parsers.plaintext import PlaintextParser
-from sumy.nlp.tokenizers import Tokenizer
-from sumy.nlp.stemmers import Stemmer
-from sumy.summarizers.lsa import LsaSummarizer
-from sumy.summarizers.luhn import LuhnSummarizer
-from sumy.utils import get_stop_words
-import itertools
-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(
-        "en_core_web_sm",
-        disable=["tok2vec", "tagger", "parser", "attribute_ruler", "lemmatizer"],
-    )
-}
-FLAIR_NER_MODELS = {"english": lambda: SequenceTagger.load("flair/ner-english")}
-REGEX_NER_MODELS = {
-    "IP_ADDRESS": [
-        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"(?:(?:\+(?:\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", "PHONE", "IP_ADDRESS", "EMAIL", "URL"]
-
-
-def _sumy__get_best_sentences(sentences, rating, *args, **kwargs):
-    from operator import attrgetter
-    from sumy.summarizers._summarizer import SentenceInfo
-
-    rate = rating
-    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 = sorted(infos, key=attrgetter("rating"), reverse=True)
-    return tuple((i.sentence, i.rating, i.order) for i in infos)
-
-
-def _sumy__lsa_call(summarizer, document):
-    summarizer._ensure_dependecies_installed()
-    dictionary = summarizer._create_dictionary(document)
-    if not dictionary:
-        return ()
-    matrix = summarizer._create_matrix(document, dictionary)
-    matrix = summarizer._compute_term_frequency(matrix)
-    from numpy.linalg import svd as singular_value_decomposition
-
-    u, sigma, v = singular_value_decomposition(matrix, full_matrices=False)
-    ranks = iter(summarizer._compute_ranks(sigma, v))
-    return _sumy__get_best_sentences(document.sentences, lambda s: next(ranks))
-
-
-def _sumy__luhn_call(summarizer, document):
-    words = summarizer._get_significant_words(document.words)
-    return _sumy__get_best_sentences(document.sentences, summarizer.rate_sentence, words)
-
-
-def get_nltk_tokenizer(language: str) -> Tokenizer:
-    nltk.data.find("tokenizers/punkt")
-    return Tokenizer(language)
-
-
-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 __repr__(self):
-        return f"NERObject(label={self.label},name={self.name})"
-
-
-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)
-    final_entities = []
-    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)
-        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,
-    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: list[NERObject] = []
-
-    # FLAIR model (if not fast)
-    if flair_model:
-        input = [Sentence(sentence[0]) for sentence in sentences]
-        flair_model.predict(input)
-        output = [e for sentence in input for e in sentence.get_spans("ner")]
-        flair_entities = [
-            NERObject(
-                name=entity.text,
-                label=BASE_TO_ONTONOTES_LABELMAP.get(
-                    entity.annotation_layers["ner"][0].value,
-                    entity.annotation_layers["ner"][0].value,
-                ),
-                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.score 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:
-            regex_entities = [
-                NERObject(
-                    name=match.group(),
-                    label=label,
-                    score=min_score - 0.5,
-                    count=1,
-                    start=match.start(),
-                )
-                for match in re.finditer(regex, text)
-            ]
-            entities += regex_entities
-    min_score = min([min_score] + [e.score for e in regex_entities])
-
-    # SPACY model
-    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 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)
-        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] + 10 * entities[i].score for i in range(len(entities))]
-        for i in range(len(entities)):
-            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]
-            comentions = [
-                entities[j].name
-                for j in range(len(entities))
-                if j != i and abs(entities[j].start - entity.start) < math.ceil(context_width / 2)
-            ]
-            entities[i].comentions = comentions
-    if with_context:
-        for i in range(len(entities)):
-            entity = entities[i]
-            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].context = context
-    return entities
-
-
-def get_extractive_summary(text, language, max_chars, fast=False, with_scores=False):
-    tokenizer = get_nltk_tokenizer(language)
-    stemmer = Stemmer(language)
-    parser = PlaintextParser.from_string(text, tokenizer)
-    if fast:
-        summarizer = LuhnSummarizer(stemmer)
-        summarizer.stop_words = get_stop_words(language)
-        scored_sentences = iter(_sumy__luhn_call(summarizer, parser.document))
-    else:
-        summarizer = LsaSummarizer(stemmer)
-        summarizer.stop_words = get_stop_words(language)
-        scored_sentences = iter(_sumy__lsa_call(summarizer, parser.document))
-    summary = []
-    summary_chars = 0
-    summary_chars_penultimate = 0
-    while summary_chars < max_chars:
-        try:
-            next_sentence = next(scored_sentences)
-            summary.append(next_sentence)
-            summary_chars_penultimate = summary_chars
-            summary_chars += len(" " + next_sentence[0]._text)
-        except StopIteration:
-            break
-    summary = sorted(summary, key=lambda x: x[2])
-    summary = [(sentence[0]._text, sentence[1]) for sentence in summary]
-    if summary_chars > max_chars:
-        summary[-1] = (
-            summary[-1][0][: max_chars - summary_chars_penultimate],
-            summary[-1][1],
-        )
-    if not with_scores:
-        summary = " ".join([s[0] for s in summary])
-    else:
-        min_score = min([s[1] for s in summary]) if summary else 0
-        max_score = max([min_score] + [s[1] for s in summary])
-        score_range = 1 if min_score == max_score else (max_score - min_score)
-        summary = [(s[0], (s[1] - min_score) / score_range) for s in summary]
-    return summary
-
-
-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, 150, with_scores, with_comentions, with_context)
-    return ner
-
-
-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, 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)