90 lines
3.0 KiB
Python
90 lines
3.0 KiB
Python
import json
|
|
import pickle
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from datasketch import MinHash, MinHashLSH
|
|
|
|
from tht.config import LshConfig
|
|
from tht.vendor.thoth_lsh import create_lsh_index, create_minhash
|
|
|
|
|
|
class LshIndexError(Exception):
|
|
pass
|
|
|
|
|
|
@dataclass
|
|
class LshHit:
|
|
table: str
|
|
column: str
|
|
value: str
|
|
score: float
|
|
|
|
|
|
def build_index(
|
|
values: dict[str, dict[str, list[str]]], cfg: LshConfig, verbose: bool = False
|
|
) -> tuple[MinHashLSH, dict[str, tuple[MinHash, str, str, str]]]:
|
|
return create_lsh_index(
|
|
values, signature_size=cfg.signature_size, n_gram=cfg.n_gram,
|
|
threshold=cfg.threshold, verbose=verbose,
|
|
)
|
|
|
|
|
|
def _paths(directory: Path, name: str) -> tuple[Path, Path, Path]:
|
|
return (
|
|
directory / f"{name}_lsh.pkl",
|
|
directory / f"{name}_minhashes.pkl",
|
|
directory / f"{name}_meta.json",
|
|
)
|
|
|
|
|
|
def save_index(lsh, minhashes, cfg: LshConfig, directory: Path, name: str) -> None:
|
|
directory.mkdir(parents=True, exist_ok=True)
|
|
lsh_path, mh_path, meta_path = _paths(directory, name)
|
|
lsh_path.write_bytes(pickle.dumps(lsh))
|
|
mh_path.write_bytes(pickle.dumps(minhashes))
|
|
meta_path.write_text(
|
|
json.dumps(
|
|
{"signature_size": cfg.signature_size, "n_gram": cfg.n_gram,
|
|
"threshold": cfg.threshold, "entries": len(minhashes)},
|
|
indent=2,
|
|
)
|
|
)
|
|
|
|
|
|
def load_index(directory: Path, name: str) -> tuple[Any, Any, dict]:
|
|
lsh_path, mh_path, meta_path = _paths(directory, name)
|
|
if not (lsh_path.exists() and mh_path.exists() and meta_path.exists()):
|
|
raise LshIndexError(
|
|
f"Indice LSH non trovato in {directory} (atteso {name}_lsh.pkl). "
|
|
f"Esegui prima `tht preprocess dwh --steps lsh`."
|
|
)
|
|
lsh = pickle.loads(lsh_path.read_bytes())
|
|
minhashes = pickle.loads(mh_path.read_bytes())
|
|
meta = json.loads(meta_path.read_text())
|
|
return lsh, minhashes, meta
|
|
|
|
|
|
def query_index(lsh, minhashes, keyword: str, meta: dict, top_n: int = 10) -> list[LshHit]:
|
|
"""Query con score: i parametri MinHash vengono dal meta dell'indice, non dalla config.
|
|
|
|
Se i due pickle (`*_lsh.pkl` e `*_minhashes.pkl`) sono disallineati (rigenerati
|
|
separatamente) `lsh.query` puo' restituire chiavi assenti da `minhashes`: si solleva
|
|
LshIndexError azionabile invece di un KeyError opaco."""
|
|
qmh = create_minhash(meta["signature_size"], keyword, meta["n_gram"])
|
|
scored = []
|
|
for key in lsh.query(qmh):
|
|
entry = minhashes.get(key)
|
|
if entry is None:
|
|
raise LshIndexError(
|
|
"Indice LSH disallineato (lsh.pkl e minhashes.pkl non coerenti): "
|
|
"rigenera con `tht preprocess dwh --steps lsh`."
|
|
)
|
|
scored.append((key, qmh.jaccard(entry[0])))
|
|
scored.sort(key=lambda kv: kv[1], reverse=True)
|
|
return [
|
|
LshHit(table=minhashes[k][1], column=minhashes[k][2], value=minhashes[k][3], score=s)
|
|
for k, s in scored[:top_n]
|
|
]
|