190 lines
6.6 KiB
Python
190 lines
6.6 KiB
Python
"""Slovak lemmatization for lexical retrieval and concept keys (Stanza).
|
|
|
|
"chatbotom", "chatbota" and "chatbot" share the lemma "chatbot"; "reči" is
|
|
"reč". Lemmas extend the full-text index and the literal keys of the activity
|
|
graph; they never replace the source text shown in answers.
|
|
|
|
The Stanza Slovak model (tokenize, pos, lemma) is prepared into the image with
|
|
the other runtime assets; no download happens at run time. Lemmas of source
|
|
units are cached on disk by text hash, so only new units are processed after
|
|
a sync. RETRIEVAL_LEMMATIZER=off (or a missing model) falls back to a rough
|
|
suffix stem; the switch exists for the documented A/B comparison.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import re
|
|
import sqlite3
|
|
import threading
|
|
from contextlib import closing
|
|
from functools import lru_cache
|
|
from pathlib import Path
|
|
|
|
from tools.answers.context_policy import normalized
|
|
|
|
LEMMA_VERSION = "stanza-sk-1"
|
|
_lock = threading.Lock()
|
|
_SK_ENDINGS = ("ami", "ach", "och", "ovi", "ych", "ymi", "eho", "emu", "ej", "om", "ov", "ou",
|
|
"ia", "ie", "iu", "a", "u", "e", "y", "i", "o")
|
|
|
|
|
|
def enabled() -> bool:
|
|
return os.getenv("RETRIEVAL_LEMMATIZER", "stanza").lower() != "off"
|
|
|
|
|
|
def word_stem(word: str) -> str:
|
|
"""Fallback: rough Slovak case stem ("chatbotom" -> "chatbot")."""
|
|
if len(word) < 5 or not word.isalpha():
|
|
return word
|
|
for ending in _SK_ENDINGS:
|
|
if word.endswith(ending) and len(word) - len(ending) >= 4:
|
|
return word[: -len(ending)]
|
|
return word
|
|
|
|
|
|
def stanza_dir() -> Path | None:
|
|
try:
|
|
from stanza.resources.common import DEFAULT_MODEL_DIR
|
|
except Exception:
|
|
DEFAULT_MODEL_DIR = None
|
|
for candidate in (os.getenv("STANZA_RESOURCES_DIR"),
|
|
os.getenv("ZP_RUNTIME_ASSETS") and str(Path(os.environ["ZP_RUNTIME_ASSETS"]) / "stanza"),
|
|
DEFAULT_MODEL_DIR):
|
|
if candidate and (Path(candidate) / "sk").is_dir():
|
|
return Path(candidate)
|
|
return None
|
|
|
|
|
|
@lru_cache(maxsize=1)
|
|
def _pipeline():
|
|
import logging
|
|
import stanza
|
|
logging.getLogger("stanza").setLevel(logging.ERROR)
|
|
directory = stanza_dir()
|
|
if directory is None:
|
|
raise RuntimeError("Stanza Slovak model is not available")
|
|
return stanza.Pipeline("sk", dir=str(directory), processors="tokenize,pos,lemma", verbose=False,
|
|
use_gpu=False, download_method=None)
|
|
|
|
|
|
def available() -> bool:
|
|
if not enabled():
|
|
return False
|
|
try:
|
|
_pipeline()
|
|
return True
|
|
except Exception:
|
|
return False
|
|
|
|
|
|
def _fallback(text: str) -> list[str]:
|
|
return [word_stem(w) for w in re.findall(r"[a-z0-9+#]+", normalized(text))]
|
|
|
|
|
|
def _from_doc(doc) -> list[str]:
|
|
out = []
|
|
for sentence in doc.sentences:
|
|
for word in sentence.words:
|
|
lemma = normalized(word.lemma or word.text)
|
|
out.extend(re.findall(r"[a-z0-9+#]+", lemma))
|
|
return out
|
|
|
|
|
|
def lemmas(text: str) -> list[str]:
|
|
"""Normalized lemma tokens (diacritics removed, lower case)."""
|
|
if not text or not text.strip():
|
|
return []
|
|
if not available():
|
|
return _fallback(text)
|
|
with _lock:
|
|
return _from_doc(_pipeline()(text))
|
|
|
|
|
|
def lemma_map(text: str) -> dict:
|
|
"""Normalized surface word -> normalized lemma for the words of a text."""
|
|
if not text or not text.strip() or not available():
|
|
return {}
|
|
with _lock:
|
|
doc = _pipeline()(text)
|
|
out = {}
|
|
for sentence in doc.sentences:
|
|
for word in sentence.words:
|
|
surface, lemma = normalized(word.text), normalized(word.lemma or word.text)
|
|
if re.fullmatch(r"[a-z0-9+#]+", surface) and re.fullmatch(r"[a-z0-9+#]+", lemma):
|
|
out[surface] = lemma
|
|
return out
|
|
|
|
|
|
def lemma_text(text: str) -> str:
|
|
return " ".join(lemmas(text))
|
|
|
|
|
|
def _key(text: str) -> str:
|
|
return hashlib.sha256((LEMMA_VERSION + "\0" + text).encode()).hexdigest()
|
|
|
|
|
|
class LemmaCache:
|
|
"""Persistent text -> lemma string cache (SQLite, one row per text hash)."""
|
|
|
|
def __init__(self, path: Path):
|
|
self.path = Path(path)
|
|
|
|
def _connect(self):
|
|
self.path.parent.mkdir(parents=True, exist_ok=True)
|
|
connection = sqlite3.connect(self.path, timeout=30)
|
|
connection.execute("CREATE TABLE IF NOT EXISTS lemmas (key TEXT PRIMARY KEY, lemmas TEXT NOT NULL)")
|
|
return connection
|
|
|
|
def many(self, texts: list[str]) -> list[str]:
|
|
"""Lemma strings for texts; missing ones are computed in bulk and stored."""
|
|
if not available():
|
|
return [" ".join(_fallback(t)) for t in texts]
|
|
keys = [_key(t) for t in texts]
|
|
found = {}
|
|
try:
|
|
with closing(self._connect()) as connection:
|
|
for start in range(0, len(keys), 500):
|
|
chunk = keys[start:start + 500]
|
|
rows = connection.execute(
|
|
f"SELECT key, lemmas FROM lemmas WHERE key IN ({','.join('?' * len(chunk))})", chunk)
|
|
found.update(rows.fetchall())
|
|
except sqlite3.Error:
|
|
pass # a cache is an optimization, never the source of truth
|
|
missing = [(k, t) for k, t in zip(keys, texts) if k not in found]
|
|
if missing:
|
|
import stanza
|
|
with _lock:
|
|
pipeline = _pipeline()
|
|
docs = pipeline.bulk_process([stanza.Document([], text=t) for _, t in missing])
|
|
fresh = {k: " ".join(_from_doc(doc)) for (k, _), doc in zip(missing, docs)}
|
|
found.update(fresh)
|
|
try:
|
|
with closing(self._connect()) as connection, connection:
|
|
connection.executemany("INSERT OR REPLACE INTO lemmas VALUES (?,?)", list(fresh.items()))
|
|
except sqlite3.Error:
|
|
pass
|
|
return [found[k] for k in keys]
|
|
|
|
|
|
def cache_path(data_dir: Path | None = None) -> Path:
|
|
from tools.rag.common import DATA_DIR
|
|
return Path(data_dir or DATA_DIR) / "knowledge" / "lemma_cache.sqlite"
|
|
|
|
|
|
def warm(data_dir: Path | None = None) -> dict:
|
|
"""Precompute lemmas of every published content unit (after a build or sync)."""
|
|
from tools.knowledge.store import load_index
|
|
from tools.sync.index_state import read_version
|
|
from tools.rag.common import DATA_DIR
|
|
data_dir = Path(data_dir or DATA_DIR)
|
|
index = load_index(data_dir, read_version(data_dir))
|
|
texts = [u["text"] for u in index["units"]]
|
|
LemmaCache(cache_path(data_dir)).many(texts)
|
|
return {"units": len(texts), "lemmatizer": LEMMA_VERSION if available() else "fallback-stem"}
|
|
|
|
|
|
if __name__ == "__main__":
|
|
print(json.dumps(warm()))
|