"""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, )