268 lines
11 KiB
Python
268 lines
11 KiB
Python
"""CPU 친화적인 영속 후보 검색 인덱스.
|
|
|
|
HashingVectorizer를 사용해 학습 vocabulary를 메모리에 들고 있지 않으며, sparse
|
|
행렬을 디스크에 저장한다. 기존 ID가 모두 유지된 경우 새 세그먼트만 transform하여
|
|
append하므로 문서 한 건 추가 시 전체 임베딩/인덱스를 재계산하지 않는다.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import hashlib
|
|
from dataclasses import dataclass
|
|
from difflib import SequenceMatcher
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
|
|
from app.engine.provenance import CorpusStore, SegmentRecord
|
|
|
|
|
|
INDEX_VERSION = 1
|
|
VECTORIZER_CONFIG = {
|
|
"analyzer": "char_wb",
|
|
"ngram_range": [3, 4],
|
|
"alternate_sign": False,
|
|
"lowercase": False,
|
|
"norm": "l2",
|
|
}
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class PersistentHit:
|
|
segment_id: str
|
|
document_id: str
|
|
title: str
|
|
score: float
|
|
evidence: list[dict]
|
|
coverage: float
|
|
longest_span: int
|
|
source_locator: str | None
|
|
coordinate_scope: str
|
|
page_number: int | None
|
|
paragraph_number: int | None
|
|
source_char_start: int | None
|
|
source_char_end: int | None
|
|
reference_text: str
|
|
|
|
|
|
def _vectorizer(n_features: int, config: dict | None = None):
|
|
from sklearn.feature_extraction.text import HashingVectorizer
|
|
|
|
cfg = config or VECTORIZER_CONFIG
|
|
return HashingVectorizer(
|
|
analyzer=cfg["analyzer"],
|
|
ngram_range=tuple(cfg["ngram_range"]),
|
|
n_features=n_features,
|
|
alternate_sign=cfg["alternate_sign"],
|
|
lowercase=cfg["lowercase"],
|
|
norm=cfg["norm"],
|
|
dtype=np.float32,
|
|
)
|
|
|
|
|
|
def _evidence_spans(query: str, reference: str, min_match: int = 12, limit: int = 10) -> tuple[list[dict], float, int]:
|
|
"""원문 query 좌표의 공통 연속 구간과 coverage를 반환."""
|
|
if not query or not reference:
|
|
return [], 0.0, 0
|
|
blocks = SequenceMatcher(None, query, reference, autojunk=False).get_matching_blocks()
|
|
useful = [b for b in blocks if b.size >= min_match]
|
|
useful.sort(key=lambda b: (-b.size, b.a))
|
|
selected = sorted(useful[:limit], key=lambda b: b.a)
|
|
spans = [
|
|
{
|
|
"start": b.a,
|
|
"end": b.a + b.size,
|
|
"source_start": b.b,
|
|
"source_end": b.b + b.size,
|
|
"matched": query[b.a : b.a + b.size],
|
|
}
|
|
for b in selected
|
|
]
|
|
covered: set[int] = set()
|
|
for b in useful:
|
|
covered.update(range(b.a, b.a + b.size))
|
|
return spans, len(covered) / max(1, len(query)), max((b.size for b in useful), default=0)
|
|
|
|
|
|
class PersistentCorpusIndex:
|
|
MATRIX_FILE = "lexical.npz"
|
|
META_FILE = "index.json"
|
|
|
|
def __init__(self, store_path: str | Path, index_dir: str | Path):
|
|
self.store = CorpusStore(store_path)
|
|
self.index_dir = Path(index_dir)
|
|
self._matrix = None
|
|
self._meta: dict = {}
|
|
|
|
@property
|
|
def ready(self) -> bool:
|
|
meta_path = self.index_dir / self.META_FILE
|
|
if not meta_path.exists():
|
|
return False
|
|
try:
|
|
meta = json.loads(meta_path.read_text(encoding="utf-8"))
|
|
return self._matrix_path(meta).exists()
|
|
except (OSError, ValueError, json.JSONDecodeError):
|
|
return False
|
|
|
|
@property
|
|
def size(self) -> int:
|
|
return len(self._meta.get("segment_ids", []))
|
|
|
|
@property
|
|
def document_count(self) -> int:
|
|
if "document_count" in self._meta:
|
|
return int(self._meta["document_count"])
|
|
return self.store.document_count()
|
|
|
|
def load(self) -> "PersistentCorpusIndex":
|
|
if not self.ready:
|
|
raise FileNotFoundError(f"영속 인덱스가 없습니다: {self.index_dir}")
|
|
from scipy.sparse import load_npz
|
|
|
|
self._meta = json.loads((self.index_dir / self.META_FILE).read_text(encoding="utf-8"))
|
|
if self._meta.get("version") != INDEX_VERSION:
|
|
raise ValueError("지원하지 않는 영속 인덱스 버전")
|
|
if self._meta.get("vectorizer_config") != VECTORIZER_CONFIG:
|
|
raise ValueError("인덱스 vectorizer 설정이 현재 코드와 달라 재빌드가 필요합니다")
|
|
self._matrix = load_npz(self._matrix_path(self._meta)).tocsr()
|
|
if self._matrix.shape[0] != len(self._meta["segment_ids"]):
|
|
raise ValueError("인덱스 행과 segment_ids 개수가 다릅니다")
|
|
return self
|
|
|
|
def _matrix_path(self, meta: dict) -> Path:
|
|
return self.index_dir / str(meta.get("matrix_file") or self.MATRIX_FILE)
|
|
|
|
def sync(self, n_features: int = 2**20) -> dict:
|
|
"""DB와 동기화. 추가만 있으면 append, 삭제/변경이면 안전하게 rebuild."""
|
|
from scipy.sparse import load_npz, save_npz, vstack
|
|
|
|
self.index_dir.mkdir(parents=True, exist_ok=True)
|
|
segments = list(self.store.iter_segments())
|
|
current = {s.segment_id: s.text_sha256 for s in segments}
|
|
mode = "rebuild"
|
|
existing_ids: list[str] = []
|
|
matrix = None
|
|
if self.ready:
|
|
old = json.loads((self.index_dir / self.META_FILE).read_text(encoding="utf-8"))
|
|
old_ids = old.get("segment_ids", [])
|
|
old_hashes = old.get("text_hashes", {})
|
|
if (
|
|
old.get("vectorizer_config") == VECTORIZER_CONFIG
|
|
and int(old.get("n_features", 0)) == n_features
|
|
and all(i in current and current[i] == old_hashes.get(i) for i in old_ids)
|
|
):
|
|
existing_ids = old_ids
|
|
matrix = load_npz(self._matrix_path(old)).tocsr()
|
|
mode = "append"
|
|
|
|
by_id = {s.segment_id: s for s in segments}
|
|
existing_set = set(existing_ids)
|
|
new_ids = [s.segment_id for s in segments if s.segment_id not in existing_set]
|
|
vectorizer = _vectorizer(n_features)
|
|
if matrix is None:
|
|
existing_ids = []
|
|
new_ids = [s.segment_id for s in segments]
|
|
matrix = vectorizer.transform([by_id[i].text for i in new_ids]).tocsr()
|
|
elif new_ids:
|
|
delta = vectorizer.transform([by_id[i].text for i in new_ids]).tocsr()
|
|
matrix = vstack([matrix, delta], format="csr")
|
|
ids = existing_ids + new_ids
|
|
generation_payload = json.dumps(
|
|
{"ids": ids, "hashes": {i: current[i] for i in ids},
|
|
"config": VECTORIZER_CONFIG, "features": n_features},
|
|
sort_keys=True,
|
|
).encode("utf-8")
|
|
generation = hashlib.sha256(generation_payload).hexdigest()[:16]
|
|
matrix_file = f"lexical-{generation}.npz"
|
|
meta = {
|
|
"version": INDEX_VERSION,
|
|
"backend": "hashing-char-3-4",
|
|
"n_features": n_features,
|
|
"vectorizer_config": VECTORIZER_CONFIG,
|
|
"segment_ids": ids,
|
|
"text_hashes": {i: current[i] for i in ids},
|
|
"document_count": self.store.document_count(),
|
|
"matrix_file": matrix_file,
|
|
}
|
|
matrix_tmp = self.index_dir / f"{matrix_file}.tmp.npz"
|
|
save_npz(matrix_tmp, matrix, compressed=False)
|
|
matrix_tmp.replace(self.index_dir / matrix_file)
|
|
tmp = self.index_dir / f"{self.META_FILE}.tmp"
|
|
tmp.write_text(json.dumps(meta, ensure_ascii=False), encoding="utf-8")
|
|
tmp.replace(self.index_dir / self.META_FILE)
|
|
self._matrix, self._meta = matrix, meta
|
|
return {"mode": mode, "total": len(ids), "added": len(new_ids)}
|
|
|
|
def query(self, text: str, top_k: int = 50, min_score: float = 0.0) -> list[PersistentHit]:
|
|
if self._matrix is None:
|
|
self.load()
|
|
if not text.strip() or self._matrix is None or self._matrix.shape[0] == 0:
|
|
return []
|
|
vectorizer = _vectorizer(
|
|
int(self._meta["n_features"]), self._meta["vectorizer_config"]
|
|
)
|
|
# 긴 원고 앞부분만 보거나 전체 벡터에 부분 복사가 희석되지 않도록 query도
|
|
# 1,200자/600 stride로 나누고 각 원문 세그먼트의 최대 점수를 사용한다.
|
|
query_chunks: list[tuple[int, str]] = []
|
|
for start in range(0, len(text), 600):
|
|
chunk = text[start : start + 1200]
|
|
if chunk.strip():
|
|
query_chunks.append((start, chunk))
|
|
if start + 1200 >= len(text):
|
|
break
|
|
scores = np.zeros(self._matrix.shape[0], dtype=np.float32)
|
|
best_chunk = np.zeros(self._matrix.shape[0], dtype=np.int32)
|
|
for batch_start in range(0, len(query_chunks), 32):
|
|
batch = query_chunks[batch_start : batch_start + 32]
|
|
queries = vectorizer.transform([chunk for _, chunk in batch])
|
|
block = (self._matrix @ queries.T).toarray()
|
|
local_argmax = block.argmax(axis=1)
|
|
local_scores = block[np.arange(block.shape[0]), local_argmax]
|
|
improved = local_scores > scores
|
|
scores[improved] = local_scores[improved]
|
|
best_chunk[improved] = batch_start + local_argmax[improved]
|
|
k = min(max(1, top_k), len(scores))
|
|
indexes = np.argpartition(scores, -k)[-k:]
|
|
indexes = indexes[np.argsort(scores[indexes])[::-1]]
|
|
ids = [self._meta["segment_ids"][int(i)] for i in indexes if scores[int(i)] >= min_score]
|
|
records = self.store.get_segments(ids)
|
|
hits: list[PersistentHit] = []
|
|
for i in indexes:
|
|
score = float(scores[int(i)])
|
|
if score < min_score:
|
|
continue
|
|
segment_id = self._meta["segment_ids"][int(i)]
|
|
record = records.get(segment_id)
|
|
if not record:
|
|
continue
|
|
chunk_start, chunk_text = query_chunks[int(best_chunk[int(i)])]
|
|
evidence, _, longest = _evidence_spans(chunk_text, record.text)
|
|
for span in evidence:
|
|
span["start"] += chunk_start
|
|
span["end"] += chunk_start
|
|
covered = sum(span["end"] - span["start"] for span in evidence)
|
|
coverage = min(1.0, covered / max(1, len(text)))
|
|
hits.append(self._to_hit(record, score, evidence, coverage, longest))
|
|
return hits
|
|
|
|
@staticmethod
|
|
def _to_hit(record: SegmentRecord, score: float, evidence: list[dict], coverage: float, longest: int) -> PersistentHit:
|
|
return PersistentHit(
|
|
segment_id=record.segment_id,
|
|
document_id=record.document_id,
|
|
title=str(record.metadata.get("document_title") or record.document_id),
|
|
score=max(0.0, min(1.0, score)),
|
|
evidence=evidence,
|
|
coverage=coverage,
|
|
longest_span=longest,
|
|
source_locator=record.source_locator,
|
|
coordinate_scope=record.coordinate_scope,
|
|
page_number=record.page_number,
|
|
paragraph_number=record.paragraph_number,
|
|
source_char_start=record.char_start,
|
|
source_char_end=record.char_end,
|
|
reference_text=record.text,
|
|
)
|