"""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 PersistentQueryResult: """질의 1건의 후보와 **문서 단위 union coverage**. coverage 정의 (#3): · ``PersistentHit.coverage`` — 이 세그먼트 하나가 질의 전체 길이에서 차지하는 비율. 분모가 질의 전체라 긴 원고에서는 필연적으로 작다. · ``PersistentQueryResult.union_coverage`` — 정밀 비교 대상 후보 전체(evidence_limit 개)의 일치 구간을 **질의 좌표에서 합집합**으로 묶어 계산한 비율. 중복 구간을 두 번 세지 않는다. "이 원고의 몇 %가 등록 코퍼스와 겹치는가"에 답하는 값은 이쪽이며, 법적 위험도 판단에는 반드시 이 값을 쓴다. """ hits: list["PersistentHit"] union_coverage: float covered_chars: int query_chars: int evidence_limit: int evidence_truncated: bool document_coverage: dict[str, float] @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 #: 인덱싱 때 캐시된 참조 특징 (#5). None 이면 호출자가 계산·백필해야 한다. reference_lemmas: list[str] | None = None reference_elements: dict | None = None #: SequenceMatcher 정밀 비교를 실제로 수행했는지. False 면 evidence/coverage/ #: longest_span 은 미계산(0) 이며 score 만 의미가 있다. reranked: bool = True 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], list[tuple[int, int]], int]: """원문 query 좌표의 공통 연속 구간을 반환. 반환: (표시용 상위 span, coverage 계산용 전체 구간 [start,end), 최장 일치 길이) 두 번째 값은 union coverage 를 위해 **잘리지 않은 전체 구간**이다. 표시용은 limit 개로 줄이지만 coverage 는 전량으로 계산해야 과소평가되지 않는다. """ if not query or not reference: return [], [], 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 ] intervals = [(b.a, b.a + b.size) for b in useful] return spans, intervals, max((b.size for b in useful), default=0) def _merge_intervals(intervals: list[tuple[int, int]]) -> list[tuple[int, int]]: """겹치는 구간을 합쳐 비중복 구간 목록으로. set(range(...)) 는 긴 원고에서 메모리를 크게 먹으므로 구간 병합으로 처리한다.""" if not intervals: return [] ordered = sorted(intervals) merged = [ordered[0]] for start, end in ordered[1:]: last_start, last_end = merged[-1] if start <= last_end: merged[-1] = (last_start, max(last_end, end)) else: merged.append((start, end)) return merged def _covered_length(intervals: list[tuple[int, int]]) -> int: return sum(end - start for start, end in _merge_intervals(intervals)) class PersistentCorpusIndex: MATRIX_FILE = "lexical.npz" META_FILE = "index.json" #: 정밀 비교(SequenceMatcher) 기본 상한. 호출자가 설정값으로 덮어쓴다. DEFAULT_EVIDENCE_LIMIT = 20 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, evidence_limit: int | None = None) -> list[PersistentHit]: """후보만 필요할 때 쓰는 얇은 래퍼. union coverage 가 필요하면 search().""" return self.search(text, top_k, min_score, evidence_limit).hits def search( self, text: str, top_k: int = 50, min_score: float = 0.0, evidence_limit: int | None = None, ) -> PersistentQueryResult: """후보 검색 + 상위 evidence_limit 개에 대해서만 정밀 비교. SequenceMatcher 는 O(질의청크 × 세그먼트) 라 후보 전체에 돌리면 요청당 수십 초가 된다. 정밀 비교 대상을 evidence_limit 로 제한하는 것이 CPU 상한이며, 나머지 후보는 score 만 채워 reranked=False 로 표시한다. """ if self._matrix is None: self.load() if not text.strip() or self._matrix is None or self._matrix.shape[0] == 0: return PersistentQueryResult([], 0.0, 0, len(text), evidence_limit or 0, False, {}) 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) limit = self.DEFAULT_EVIDENCE_LIMIT if evidence_limit is None else evidence_limit limit = max(0, limit) hits: list[PersistentHit] = [] union_intervals: list[tuple[int, int]] = [] document_intervals: dict[str, list[tuple[int, int]]] = {} reranked_count = 0 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 if reranked_count >= limit: # CPU 상한. 정밀 비교 없이 score 만 채운다. hits.append(self._to_hit(record, score, [], 0.0, 0, reranked=False)) continue chunk_start, chunk_text = query_chunks[int(best_chunk[int(i)])] evidence, intervals, longest = _evidence_spans(chunk_text, record.text) for span in evidence: span["start"] += chunk_start span["end"] += chunk_start shifted = [(s + chunk_start, e + chunk_start) for s, e in intervals] union_intervals.extend(shifted) document_intervals.setdefault(record.document_id, []).extend(shifted) # 이 세그먼트 단독 기여분 (질의 전체 길이 대비) coverage = min(1.0, _covered_length(shifted) / max(1, len(text))) hits.append(self._to_hit(record, score, evidence, coverage, longest)) reranked_count += 1 covered_chars = _covered_length(union_intervals) return PersistentQueryResult( hits=hits, union_coverage=min(1.0, covered_chars / max(1, len(text))), covered_chars=covered_chars, query_chars=len(text), evidence_limit=limit, evidence_truncated=any(not h.reranked for h in hits), document_coverage={ document_id: min(1.0, _covered_length(intervals) / max(1, len(text))) for document_id, intervals in document_intervals.items() }, ) @staticmethod def _to_hit(record: SegmentRecord, score: float, evidence: list[dict], coverage: float, longest: int, reranked: bool = True) -> 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, reference_lemmas=record.lemmas, reference_elements=record.elements, reranked=reranked, )