375 lines
15 KiB
Python
375 lines
15 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 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,
|
||
)
|