File size: 9,246 Bytes
ac4a6d6 | 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 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 | import os
import logging
import spacy
from typing import Dict, Any, List
# Base backend directory
BASE_DIR = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
MODEL_DIR = os.path.join(BASE_DIR, "model")
# Define paths to custom spaCy models
model_path1 = os.path.join(MODEL_DIR, "model-best")
model_path2 = os.path.join(MODEL_DIR, "2")
model_path3 = os.path.join(MODEL_DIR, "3")
# Memory cache for lazy loading spaCy and Transformer models
_spacy_models_cache: Dict[int, Any] = {}
_transformer_pipeline = None
import time
# Default fallback cloud model
FALLBACK_TRANSFORMER_MODEL = "dslim/distilbert-NER"
def get_target_transformer_model():
"""
Check for custom trained Transformer model paths or environment variables.
Returns the path to the custom model if found, otherwise returns FALLBACK_TRANSFORMER_MODEL.
"""
env_model = os.environ.get("NER_TRANSFORMER_MODEL")
if env_model:
return env_model
# Check potential local trained model directories
possible_paths = [
os.path.join(BASE_DIR, "training", "roberta_ner_model"),
os.path.join(MODEL_DIR, "roberta_ner_model"),
os.path.join(MODEL_DIR, "roberta"),
os.path.join(BASE_DIR, "roberta_ner_model")
]
for path in possible_paths:
if os.path.exists(path):
logging.info(f"Auto-detected custom trained Transformer model at: {path}")
return path
return FALLBACK_TRANSFORMER_MODEL
_transformer_model_name = get_target_transformer_model()
def get_spacy_model(level: int = 1):
"""
Lazy load spaCy models into memory cache with graceful fallback to standard en_core_web models.
"""
if level in _spacy_models_cache:
return _spacy_models_cache[level]
path = model_path1 if level == 1 else (model_path2 if level == 2 else model_path3)
fallback = "en_core_web_sm" if level == 1 else ("en_core_web_md" if level == 2 else "en_core_web_lg")
try:
if not os.path.exists(path):
spacy_repo = os.environ.get("NER_SPACY_REPO")
if spacy_repo:
try:
from huggingface_hub import snapshot_download
logging.info(f"Downloading custom spaCy models from HuggingFace repo {spacy_repo}...")
snapshot_download(repo_id=spacy_repo, local_dir=MODEL_DIR)
except Exception as e:
logging.warning(f"Failed to download spaCy repo {spacy_repo}: {e}")
if os.path.exists(path):
logging.info(f"Lazy loading spaCy model from {path}...")
_spacy_models_cache[level] = spacy.load(path)
else:
logging.info(f"Custom model not found at {path}. Loading fallback '{fallback}'...")
try:
_spacy_models_cache[level] = spacy.load(fallback)
except OSError:
logging.warning(f"Fallback {fallback} not installed. Loading blank en model...")
_spacy_models_cache[level] = spacy.blank("en")
except Exception as e:
logging.error(f"Error loading spaCy model level {level}: {e}. Falling back to blank model.")
_spacy_models_cache[level] = spacy.blank("en")
return _spacy_models_cache[level]
def get_transformer_pipeline():
"""
Lazy load HuggingFace Transformer NER pipeline.
Tries to load the custom trained model up to 3 times before falling back to the downloaded cloud model.
"""
global _transformer_pipeline
if _transformer_pipeline is None:
try:
from transformers import pipeline
target_model = get_target_transformer_model()
models_to_try = []
# If custom model is selected, attempt to load it up to 3 times
if target_model != FALLBACK_TRANSFORMER_MODEL:
for attempt_num in range(1, 4):
models_to_try.append((target_model, attempt_num, True))
# Finally, add fallback cloud model as the last resort
models_to_try.append((FALLBACK_TRANSFORMER_MODEL, 1, False))
for model_path, attempt, is_custom in models_to_try:
try:
if is_custom:
logging.info(f"Attempting to load custom trained Transformer model from '{model_path}' (Attempt {attempt} of 3)...")
else:
logging.info(f"Loading fallback downloaded cloud model '{model_path}'...")
try:
_transformer_pipeline = pipeline("ner", model=model_path, aggregation_strategy="simple")
except TypeError:
_transformer_pipeline = pipeline("ner", model=model_path, grouped_entities=True)
logging.info(f"Transformer NER pipeline successfully loaded using: {model_path}")
break
except Exception as e:
if is_custom:
logging.warning(f"Failed to load custom model on Attempt {attempt}: {e}")
if attempt < 3:
time.sleep(1) # short pause before retry
else:
logging.warning(f"Custom trained model failed after 3 attempts! Switching to downloaded fallback model '{FALLBACK_TRANSFORMER_MODEL}'.")
else:
logging.error(f"Error loading fallback transformer model '{model_path}': {e}. Falling back to spaCy NER.")
return None
except ImportError:
logging.warning("Transformers library not installed. Transformer NER will fallback to spaCy.")
return None
except Exception as e:
logging.error(f"Unexpected error initializing Transformer pipeline: {e}. Falling back to spaCy NER.")
return None
return _transformer_pipeline
def redact_entities_spacy(text: str, level: int = 1) -> str:
"""
Unified spaCy NER redaction function.
Uses character-span slice replacement (O(n log n)) sorted descending to prevent O(n²) str.replace() clobbering.
"""
if not text or not isinstance(text, str):
return text
doc = get_spacy_model(level)(text)
entities = []
for ent in doc.ents:
if ent.text and ent.text.strip():
entities.append({
"start": ent.start_char,
"end": ent.end_char,
"text": ent.text
})
# Sort descending by start index so earlier replacements don't shift subsequent character indices
entities_sorted = sorted(entities, key=lambda x: x["start"], reverse=True)
redacted_text = text
for ent in entities_sorted:
start, end = ent["start"], ent["end"]
if start is not None and end is not None and 0 <= start < len(redacted_text) and end <= len(redacted_text):
mask = 'x' * (end - start)
redacted_text = redacted_text[:start] + mask + redacted_text[end:]
return redacted_text
def redact_entities_transformer(text: str, mask_char: str = "x") -> str:
"""
Redact named entities using deep learning Transformer model (DistilBERT/RoBERTa).
Falls back to spaCy Level 1 if transformers cannot be loaded.
"""
if not text or not isinstance(text, str):
return text
pipe = get_transformer_pipeline()
if pipe is None:
return redact_entities_spacy(text, level=1)
try:
entities = pipe(text)
entities_sorted = sorted(entities, key=lambda x: x.get('start', 0), reverse=True)
redacted_text = text
for ent in entities_sorted:
word = ent.get('word', '')
score = ent.get('score', 0.0)
if score > 0.40 and word and len(word.strip()) > 1:
start, end = ent.get('start'), ent.get('end')
if start is not None and end is not None and 0 <= start < len(redacted_text) and end <= len(redacted_text):
mask = mask_char * (end - start)
redacted_text = redacted_text[:start] + mask + redacted_text[end:]
return redacted_text
except Exception as e:
logging.error(f"Error in transformer redaction: {e}")
return redact_entities_spacy(text, level=1)
def preload_fallback_model():
"""
Pre-download / cache the fallback cloud model at startup so there is zero download wait time during fallback.
"""
try:
from transformers import pipeline
logging.info(f"Pre-downloading/caching fallback Transformer model '{FALLBACK_TRANSFORMER_MODEL}' at startup...")
try:
pipeline("ner", model=FALLBACK_TRANSFORMER_MODEL, aggregation_strategy="simple")
except TypeError:
pipeline("ner", model=FALLBACK_TRANSFORMER_MODEL, grouped_entities=True)
logging.info("Fallback Transformer model successfully pre-cached at startup!")
except Exception as e:
logging.warning(f"Could not pre-cache fallback transformer model at startup: {e}")
|