Coverage for src/asketmc_bot/lemma.py: 0%

1#!/usr/bin/env python3.10

2"""

3lemma.py — lemmatizer (stanza/spaCy) and lemma caches for files/chunks.

5Exported API:

6- extract_lemmas(text: str) -> FrozenSet[str]

7- FILE_LEMMAS: Dict[str, FrozenSet[str]]

8- CHUNK_LEMMA_CACHE: Dict[str, FrozenSet[str]]

9- CHUNK_LEMMA_CACHE_FILE: Path

10- LEMMA_INDEX_FILE: Path

11- save_chunk_lemma_cache(...)

12- load_saved_lemmas()

13- update_file_lemmas_async(docs, stored_hashes, new_hashes)

14- LEMMA_POOL (ThreadPoolExecutor) — for controlled shutdown if needed

16Notes:

17- No model downloads or heavy initialization at import time.

18- Model init is lazy. Optional stanza resources auto-download can happen on first use.

19- Cache files are stored under cfg.CACHE_PATH.

20"""

22from __future__ import annotations

24# Stdlib

25import asyncio

26import functools

27import hashlib

28import json

29import logging

30import os

31import threading

32import time

33from concurrent.futures import ThreadPoolExecutor

34from pathlib import Path

35from typing import Dict, FrozenSet, List, Optional, Tuple

37# Third-party

38import spacy

39import stanza

40from langdetect import detect

41from langdetect.lang_detect_exception import LangDetectException

43# Local

44from asketmc_bot import config as cfg

46log = logging.getLogger("asketmc.lemma")

47log.setLevel(logging.DEBUG if getattr(cfg, "DEBUG", False) else logging.INFO)

49# ─────────────────────────────────────────────────────────

50# Caches (stored under cfg.CACHE_PATH)

51# ─────────────────────────────────────────────────────────

52CHUNK_LEMMA_CACHE_FILE: Path = cfg.CACHE_PATH / "chunk_lemma_index.json"

53LEMMA_INDEX_FILE: Path = cfg.CACHE_PATH / "lemma_index.json"

55CHUNK_LEMMA_CACHE: Dict[str, FrozenSet[str]] = {}

56FILE_LEMMAS: Dict[str, FrozenSet[str]] = {}

58# ─────────────────────────────────────────────────────────

59# Pool and locks

60# ─────────────────────────────────────────────────────────

61_LEMMA_LOCK = threading.Lock()

62_MODELS_INIT_LOCK = threading.Lock()

63LEMMA_POOL = ThreadPoolExecutor(max_workers=min(os.cpu_count() or 4, 8))

65# ─────────────────────────────────────────────────────────

66# Lazy model initialization

67# ─────────────────────────────────────────────────────────

68_STANZA_NLP_RU: Optional[stanza.Pipeline] = None

69_SPACY_EN = None

71# Tracks readiness separately (important for cache invalidation).

72_SPACY_READY = False

73_STANZA_READY = False

75_MODELS_INIT_TRIED = False

76_LAST_INIT_ATTEMPT_TS = 0.0

77_INIT_RETRY_COOLDOWN_SEC = float(getattr(cfg, "LEMMA_INIT_RETRY_COOLDOWN_SEC", 300.0))

79# Degraded-mode reporting (one-time)

80_DEGRADED_WARNED = False

82# Cache epoch: changes whenever model availability changes.

83# Used to prevent "empty lemma" results from being cached forever.

84_MODELS_EPOCH = 0

86# Stanza resources location + auto-download switch (lazy, on first use).

87_STANZA_DIR: Path = Path(

88 getattr(cfg, "STANZA_RESOURCES_DIR", cfg.CACHE_PATH / "stanza")

89).resolve()

90_STANZA_AUTO_DOWNLOAD: bool = bool(getattr(cfg, "STANZA_AUTO_DOWNLOAD", True))

93def _invalidate_all_lemma_caches(reason: str) -> None:

94 """Clear all lemma-related caches and persisted cache files."""

95 global _MODELS_EPOCH

96 _MODELS_EPOCH += 1

98 try:

99 _extract_lemmas_cached.cache_clear()

100 except Exception:

101 pass

102

103 CHUNK_LEMMA_CACHE.clear()

104 FILE_LEMMAS.clear()

105

106 for fp in (LEMMA_INDEX_FILE, CHUNK_LEMMA_CACHE_FILE):

107 try:

108 if fp.exists():

109 fp.unlink()

110 except Exception:

111 pass

112

113 log.info(

114 "[LEMMA_CACHE] invalidated (epoch=%d) reason=%s",

115 _MODELS_EPOCH,

116 reason,

117 )

118

119

120def _notify_degraded_once() -> None:

121 global _DEGRADED_WARNED

122 if _DEGRADED_WARNED:

123 return

124 _DEGRADED_WARNED = True

125 log.warning(

126 "[LEMMATIZER] Lemmatization is in degraded mode: spaCy EN model and/or stanza RU resources are unavailable. "

127 "Lemmas may be empty until models are available."

128 )

129

130

131def _lang_of(text: str) -> str:

132 """Best-effort language detection. Returns 'en' or 'ru'."""

133 try:

134 detected_lang = detect(text)

135 except LangDetectException:

136 detected_lang = "ru"

137 except Exception:

138 detected_lang = "ru"

139 return "en" if detected_lang == "en" else "ru"

140

141

142def _looks_like_torch_weights_only_error(err: Exception) -> bool:

143 """Heuristic to detect PyTorch 2.6+ 'weights_only' compatibility failures surfaced via stanza."""

144 msg = str(err)

145 needles = (

146 "Weights only load failed",

147 "WeightsUnpickler error",

148 "torch.serialization.add_safe_globals",

149 "weights_only",

150 "Unsupported global",

151 )

152 return any(n in msg for n in needles)

153

154

155def _torch_allowlist_stanza_safe_globals() -> bool:

156 """

157 Allowlist numpy globals required by torch weights-only unpickler (PyTorch 2.6+).

158 Only use if checkpoints are trusted.

159 """

160 try:

161 import numpy as np # local import

162 import torch # local import

163 except Exception as exc:

164 log.warning("[LEMMATIZER] torch/numpy import failed: %s", exc)

165 return False

166

167 ser = getattr(torch, "serialization", None)

168 add = getattr(ser, "add_safe_globals", None) if ser else None

169 if not callable(add):

170 log.warning(

171 "[LEMMATIZER] torch.serialization.add_safe_globals not available; cannot relax weights_only safely"

172 )

173 return False

174

175 try:

176 allow = [np.core.multiarray._reconstruct, np.ndarray, np.dtype]

177 if hasattr(np.core.multiarray, "scalar"):

178 allow.append(np.core.multiarray.scalar)

179 add(allow)

180 return True

181 except Exception as exc:

182 log.warning("[LEMMATIZER] torch safe-globals allowlist failed: %s", exc)

183 return False

184

185

186class _TorchLoadWeightsOnlyFalse:

187 """

188 Narrow monkeypatch: force torch.load(weights_only=False) during stanza pipeline init.

189 Only use if checkpoints are trusted.

190 """

191

192 def __enter__(self):

193 import torch # local import

194

195 self._torch = torch

196 self._orig = torch.load

197

198 def _patched(*args, **kwargs):

199 kwargs.setdefault("weights_only", False)

200 return self._orig(*args, **kwargs)

201

202 torch.load = _patched

203 return self

204

205 def __exit__(self, exc_type, exc, tb):

206 self._torch.load = self._orig

207 return False

208

209

210def _try_init_models(*, force_retry: bool = False) -> Tuple[bool, bool]:

211 """

212 Lazily initialize spaCy EN + stanza RU.

213

214 Returns (spacy_ready, stanza_ready).

215 May trigger stanza resource download (lazy) if STANZA_AUTO_DOWNLOAD=True.

216

217 Thread-safety:

218 - Single initializer under _MODELS_INIT_LOCK.

219

220 Retry behavior:

221 - If init previously failed, allow throttled retry (cooldown) or explicit force_retry.

222 """

223 global _STANZA_NLP_RU, _SPACY_EN, _MODELS_INIT_TRIED, _LAST_INIT_ATTEMPT_TS

224 global _SPACY_READY, _STANZA_READY

225

226 now = time.time()

227

228 if (

229 _MODELS_INIT_TRIED

230 and (now - _LAST_INIT_ATTEMPT_TS) < _INIT_RETRY_COOLDOWN_SEC

231 and not force_retry

232 ):

233 return _SPACY_READY, _STANZA_READY

234

235 with _MODELS_INIT_LOCK:

236 now = time.time()

237 if (

238 _MODELS_INIT_TRIED

239 and (now - _LAST_INIT_ATTEMPT_TS) < _INIT_RETRY_COOLDOWN_SEC

240 and not force_retry

241 ):

242 return _SPACY_READY, _STANZA_READY

243

244 _LAST_INIT_ATTEMPT_TS = now

245 _MODELS_INIT_TRIED = True

246

247 prev_spacy = _SPACY_READY

248 prev_stanza = _STANZA_READY

249

250 # ---- spaCy EN ----

251 spacy_en = None

252 try:

253 spacy_en = spacy.load("en_core_web_sm")

254 _SPACY_READY = True

255 log.info("[LEMMATIZER] spaCy en_core_web_sm ready")

256 except Exception as e:

257 _SPACY_READY = False

258 log.warning("[LEMMATIZER] spaCy model not available: %s", e)

259

260 # ---- stanza RU ----

261 stanza_ru = None

262 _STANZA_DIR.mkdir(parents=True, exist_ok=True)

263

264 def _build_stanza_pipeline() -> stanza.Pipeline:

265 return stanza.Pipeline(

266 lang="ru",

267 processors="tokenize,pos,lemma",

268 use_gpu=False,

269 verbose=False,

270 dir=str(_STANZA_DIR),

271 )

272

273 try:

274 stanza_ru = _build_stanza_pipeline()

275 _STANZA_READY = True

276 log.info("[LEMMATIZER] stanza ru pipeline ready (dir=%s)", _STANZA_DIR)

277 except Exception as e:

278 _STANZA_READY = False

279

280 if _looks_like_torch_weights_only_error(e):

281 log.warning(

282 "[LEMMATIZER] stanza ru pipeline blocked by torch serialization restrictions; "

283 "trying safe-globals allowlist (dir=%s). Error=%s",

284 _STANZA_DIR,

285 e,

286 )

287

288 if _torch_allowlist_stanza_safe_globals():

289 try:

290 stanza_ru = _build_stanza_pipeline()

291 _STANZA_READY = True

292 log.info(

293 "[LEMMATIZER] stanza ru pipeline ready after torch safe-globals allowlist (dir=%s)",

294 _STANZA_DIR,

295 )

296 except Exception as e2:

297 _STANZA_READY = False

298 log.warning(

299 "[LEMMATIZER] stanza still blocked after safe-globals allowlist; "

300 "trying torch.load(weights_only=False) (dir=%s). Error=%s",

301 _STANZA_DIR,

302 e2,

303 )

304

305 if not _STANZA_READY:

306 try:

307 with _TorchLoadWeightsOnlyFalse():

308 stanza_ru = _build_stanza_pipeline()

309 _STANZA_READY = True

310 log.info(

311 "[LEMMATIZER] stanza ru pipeline ready after forcing torch.load(weights_only=False) (dir=%s)",

312 _STANZA_DIR,

313 )

314 except Exception as e3:

315 _STANZA_READY = False

316 log.warning(

317 "[LEMMATIZER] stanza ru pipeline blocked even after weights_only=False; "

318 "stanza disabled (dir=%s). Error=%s",

319 _STANZA_DIR,

320 e3,

321 )

322 else:

323 log.warning("[LEMMATIZER] stanza pipeline not available: %s", e)

324

325 if _STANZA_AUTO_DOWNLOAD:

326 try:

327 log.info(

328 "[LEMMATIZER] downloading stanza resources for ru (model_dir=%s)",

329 _STANZA_DIR,

330 )

331 stanza.download(

332 "ru",

333 model_dir=str(_STANZA_DIR),

334 processors="tokenize,pos,lemma",

335 verbose=False,

336 )

337 stanza_ru = _build_stanza_pipeline()

338 _STANZA_READY = True

339 log.info(

340 "[LEMMATIZER] stanza ru pipeline ready after download (dir=%s)",

341 _STANZA_DIR,

342 )

343 except Exception as e2:

344 _STANZA_READY = False

345 log.warning("[LEMMATIZER] stanza download/init failed: %s", e2)

346

347 _SPACY_EN = spacy_en

348 _STANZA_NLP_RU = stanza_ru

349

350 if (_SPACY_READY and not prev_spacy) or (_STANZA_READY and not prev_stanza):

351 _invalidate_all_lemma_caches(

352 reason=f"models_became_available spacy={_SPACY_READY} stanza={_STANZA_READY}"

353 )

354

355 return _SPACY_READY, _STANZA_READY

356

357

358@functools.lru_cache(maxsize=10_000)

359def _extract_lemmas_cached(text: str, epoch: int) -> FrozenSet[str]:

360 """Cached lemmatization; epoch prevents degraded empty results from sticking forever."""

361 _ = epoch # part of cache key; not used directly

362

363 if not text or len(text) < 3:

364 return frozenset()

365

366 spacy_ready, stanza_ready = _try_init_models(force_retry=False)

367 if not spacy_ready and not stanza_ready:

368 _notify_degraded_once()

369 return frozenset()

370

371 lang = _lang_of(text)

372

373 if lang == "en":

374 if not spacy_ready or _SPACY_EN is None:

375 return frozenset()

376 try:

377 doc = _SPACY_EN(text)

378 lemmas = {

379 tok.lemma_.lower()

380 for tok in doc

381 if tok.is_alpha and not tok.is_stop and len(tok.text) > 2

382 }

383 return frozenset(lemmas)

384 except Exception:

385 return frozenset()

386

387 if not stanza_ready or _STANZA_NLP_RU is None:

388 return frozenset()

389

390 try:

391 with _LEMMA_LOCK:

392 doc = _STANZA_NLP_RU(text)

393 except Exception:

394 return frozenset()

395

396 try:

397 lemmas = {

398 w.lemma.lower()

399 for s in doc.sentences

400 for w in s.words

401 if w.lemma

402 and len(w.lemma) > 2

403 and w.upos in cfg.GOOD_POS

404 and w.lemma.lower() not in cfg.STOP_WORDS

405 }

406 return frozenset(lemmas)

407 except Exception:

408 return frozenset()

409

410

411def extract_lemmas(text: str) -> FrozenSet[str]:

412 """Public API: lemmatize text with epoch-aware caching."""

413 return _extract_lemmas_cached(text, _MODELS_EPOCH)

414

415

416def _chunk_hash(text: str) -> str:

417 return hashlib.sha256(text.encode("utf-8")).hexdigest()

418

419

420def get_lemmas_for_chunk(text: str) -> FrozenSet[str]:

421 h = _chunk_hash(text)

422 lemmas = CHUNK_LEMMA_CACHE.get(h)

423 if lemmas is None:

424 try:

425 lemmas = extract_lemmas(text)

426 except Exception:

427 lemmas = frozenset()

428 CHUNK_LEMMA_CACHE[h] = lemmas

429 return lemmas

430

431

432def save_chunk_lemma_cache(chunk_cache_file: Path = CHUNK_LEMMA_CACHE_FILE) -> None:

433 data = {k: list(v) for k, v in CHUNK_LEMMA_CACHE.items()}

434 try:

435 chunk_cache_file.parent.mkdir(parents=True, exist_ok=True)

436 chunk_cache_file.write_text(

437 json.dumps(data, ensure_ascii=False, indent=2),

438 encoding="utf-8",

439 )

440 log.info("[LEMMA_CACHE] Chunk cache saved: %s (keys=%d)", chunk_cache_file, len(data))

441 except Exception as e:

442 log.error("[LEMMA_CACHE] Save error: %s", e)

443

444

445def _persist_lemmas(lemma_file: Optional[Path] = None) -> None:

446 lemma_file = lemma_file or LEMMA_INDEX_FILE

447 try:

448 lemma_file.parent.mkdir(parents=True, exist_ok=True)

449 dump = {k: list(v) for k, v in FILE_LEMMAS.items()}

450 lemma_file.write_text(

451 json.dumps(dump, ensure_ascii=False, indent=2),

452 encoding="utf-8",

453 )

454 log.info("[LEMMA_CACHE] File-lemmas saved: %s (files=%d)", lemma_file, len(FILE_LEMMAS))

455 except Exception as e:

456 log.error("[LEMMA_CACHE] Persist lemmas error: %s", e)

457

458

459def load_saved_lemmas() -> None:

460 if not LEMMA_INDEX_FILE.exists():

461 log.info("[LEMMA_CACHE] Lemma cache not found: %s", LEMMA_INDEX_FILE)

462 return

463 try:

464 data = json.loads(LEMMA_INDEX_FILE.read_text("utf-8"))

465 if isinstance(data, dict):

466 for fname, lst in data.items():

467 if isinstance(lst, list):

468 FILE_LEMMAS[fname] = frozenset(lst)

469 log.info("[LEMMA_CACHE] Loaded file-lemmas: %d files", len(FILE_LEMMAS))

470 except Exception as e:

471 log.warning("[LEMMA_CACHE] Load lemmas error: %s", e)

472

473

474def _read_file(fp: Path) -> str:

475 try:

476 return fp.read_text("utf-8", "ignore")

477 except Exception:

478 return ""

479

480

481async def _compute_and_store_lemmas(fp: Path) -> None:

482 loop = asyncio.get_running_loop()

483 text = await loop.run_in_executor(LEMMA_POOL, _read_file, fp)

484 try:

485 FILE_LEMMAS[fp.name] = extract_lemmas(text)

486 log.debug("[LEMMA_CACHE] %s: %d lemmas", fp.name, len(FILE_LEMMAS[fp.name]))

487 except Exception:

488 FILE_LEMMAS[fp.name] = frozenset()

489 log.debug("[LEMMA_CACHE] %s: set empty (error)", fp.name)

490

491

492async def update_file_lemmas_async(

493 docs: List[Path],

494 stored_hashes: Dict[str, str],

495 new_hashes: Dict[str, str],

496) -> List[Path]:

497 changed = [d for d in docs if stored_hashes.get(d.name) != new_hashes.get(d.name)]

498 if not changed:

499 log.info("[LEMMA_CACHE] No changed files")

500 return []

501

502 tasks = [asyncio.create_task(_compute_and_store_lemmas(d)) for d in changed]

503 await asyncio.gather(*tasks)

504 _persist_lemmas()

505 log.info("[LEMMA_CACHE] Updated file-lemmas for %d files", len(changed))

506 return changed