Spaces:
Running
Running
File size: 4,739 Bytes
a55ea57 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 | """
File: /ner.py
Project: nlp
Created Date: Sunday February 8th 2026
Author: Christian Nonis <alch.infoemail@gmail.com>
-----
Last Modified: Sunday February 8th 2026 11:21:15 am
Modified By: Christian Nonis <alch.infoemail@gmail.com>
-----
"""
from pydantic import BaseModel
from src.lib.tracing.profiler import profile_stage
from src.utils.nlp.spacy import MODEL_NAMES, _spacy_manager
from src.utils.nlp.lang_detect import langid_detect
_DEFAULT_IGNORE_POS = frozenset(
{
"PUNCT",
"SPACE",
"DET",
"PRON",
"AUX",
"CCONJ",
"ADP",
"SCONJ",
"PART",
"NUM",
"SYM",
"X",
}
)
class ExtractElementsResponse(BaseModel):
tokens: list[dict[str, str]]
noun_chunks: list[str]
class MultiLangEntityExtractor:
def __init__(self):
self.spacy_manager = _spacy_manager
def extract(self, text):
lang, _ = langid_detect(text)
if lang not in MODEL_NAMES:
return []
try:
nlp = self.spacy_manager.get_model(lang)
except Exception:
return []
doc = nlp(text)
return [(ent.text, ent.label_) for ent in doc.ents]
def _process_doc(self, doc, ignore_pos_set):
tokens = []
i = 0
n = len(doc)
while i < n:
t = doc[i]
if t.is_space:
i += 1
continue
if t.ent_iob_ == "B":
span = [t]
j = 1
while i + j < n and doc[i + j].ent_iob_ == "I":
span.append(doc[i + j])
j += 1
tokens.append(
{
"text": " ".join(x.text for x in span),
"lemma": " ".join(x.lemma_ for x in span),
"pos": span[0].pos_,
"entity": t.ent_type_,
}
)
i += len(span)
continue
if t.ent_iob_ == "I":
i += 1
continue
if t.pos_ not in ignore_pos_set:
tokens.append({"text": t.text, "lemma": t.lemma_, "pos": t.pos_})
i += 1
entity_spans = [(ent.start, ent.end) for ent in doc.ents]
noun_chunks = []
for chunk in doc.noun_chunks:
for start, end in entity_spans:
if start <= chunk.start and chunk.end <= end:
break
else:
noun_chunks.append(chunk.text.strip())
for ent in doc.ents:
noun_chunks.append(ent.text.strip())
return {"tokens": tokens, "noun_chunks": noun_chunks}
def extract_elements(
self,
text,
lang: str | None = None,
ignore_pos: list[str] | None = None,
) -> ExtractElementsResponse:
if lang is None:
with profile_stage("nlp.detect_language"):
lang, _ = langid_detect(text)
if lang not in MODEL_NAMES:
return {"tokens": [], "noun_chunks": []}
try:
with profile_stage("nlp.load_model", lang=lang):
nlp = self.spacy_manager.get_model(lang)
except Exception:
return {"tokens": [], "noun_chunks": []}
with profile_stage("nlp.parse", chars=len(text or "")):
doc = nlp(text)
ignore_pos_set = (
set(ignore_pos) if ignore_pos is not None else _DEFAULT_IGNORE_POS
)
with profile_stage("nlp.process_doc"):
tokens = self._process_doc(doc, ignore_pos_set)["tokens"]
with profile_stage("nlp.process_doc"):
noun_chunks = self._process_doc(doc, ignore_pos_set)["noun_chunks"]
return ExtractElementsResponse(tokens=tokens, noun_chunks=noun_chunks)
def extract_elements_batch(
self,
texts: list[str],
lang: str | None = None,
ignore_pos: list[str] | None = None,
batch_size: int = 50,
):
if not texts:
return []
if lang is None:
lang, _ = langid_detect(texts[0])
if lang not in MODEL_NAMES:
return [{"tokens": [], "noun_chunks": []} for _ in texts]
try:
nlp = self.spacy_manager.get_model(lang)
except Exception:
return [{"tokens": [], "noun_chunks": []} for _ in texts]
ignore_pos_set = (
set(ignore_pos) if ignore_pos is not None else _DEFAULT_IGNORE_POS
)
out = []
for doc in nlp.pipe(texts, batch_size=batch_size):
out.append(self._process_doc(doc, ignore_pos_set))
return out
_entity_extractor = MultiLangEntityExtractor()
|