Spaces:
Running
Running
File size: 2,013 Bytes
fad436e | 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 | import logging
from typing import List, Dict
from .base import BaseNER
class NEREngine(BaseNER):
def __init__(self, model_name="urchade/gliner_mediumv2.1"):
self.model_name = model_name
self.model = None
self._initialize_model()
def _initialize_model(self):
logging.info(f"Initializing NER model: {self.model_name}")
try:
from backup.model import GLiNER
self.model = GLiNER.from_pretrained(self.model_name)
logging.info(f"NER model '{self.model_name}' loaded successfully.")
except Exception as e:
logging.error(f"Failed to load NER model: {e}. NER extraction will be unavailable.")
def extract_entities(self, text: str, labels: List[str] = None) -> Dict[str, List[str]]:
if not text:
logging.warning("NER: Received empty text for extraction.")
return {}
if not self.model:
logging.error("NER: Model not initialized. Skipping extraction.")
return {}
if labels is None:
labels = ["Name", "Designation", "Company", "Contact", "Address", "Email", "Link"]
logging.info(f"NER: Extracting entities for {len(text)} characters of text.")
try:
entities = self.model.predict_entities(text, labels, threshold=0.3)
structured_data = {label: [] for label in labels}
for ent in entities:
label = ent["label"]
if label in structured_data:
structured_data[label].append(ent["text"])
non_empty_tags = sum(1 for v in structured_data.values() if v)
logging.info(f"NER: Extracted entities for {non_empty_tags} labels.")
return structured_data
except Exception as e:
logging.error(f"NER: Extraction pipeline crashed: {e}")
return {}
def process(self, text: str) -> Dict[str, List[str]]:
return self.extract_entities(text)
|