Spaces:
Sleeping
Sleeping
File size: 1,901 Bytes
8db761b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 | from __future__ import annotations
import re
from dataclasses import dataclass
import numpy as np
from ..ingest.embedder import Embedder
from ..ingest.models import Chunk
from .loader import RetrievalIndex
@dataclass
class ScoredChunk:
chunk: Chunk
score: float
def _norm(s: str) -> str:
return re.sub(r"\s+", " ", s.lower()).strip()
def retrieve(query: str, index: RetrievalIndex, embedder: Embedder,
k: int = 6, candidate_k: int = 30, reranker=None) -> list[ScoredChunk]:
n = len(index.chunks)
if n == 0:
return []
qv = embedder.encode([query]).astype("float32") # (1, dim), L2-normalized
rows: set[int] = set()
_, ids = index.faiss_index.search(qv, min(candidate_k, n))
for r in ids[0]:
if r >= 0:
rows.add(int(r))
# lexical concept-expansion (centroid vectors are deferred to Phase 3)
nq = _norm(query)
for concept, chunk_ids in index.concept_index.items():
cn = _norm(concept)
if cn and cn in nq:
for cid in chunk_ids:
r = index.row_by_id.get(cid)
if r is not None:
rows.add(r)
q = qv[0]
scored = [
ScoredChunk(chunk=index.chunks[r],
score=float(np.dot(q, np.asarray(index.faiss_index.reconstruct(int(r)), dtype="float32"))))
for r in rows
]
scored.sort(key=lambda s: s.score, reverse=True)
# Optional rerank: a cross-encoder re-scores the top cosine candidates for sharper
# grounding. Cap the pool so CPU latency stays bounded (near-instant on GPU).
if reranker is not None and scored:
pool = scored[: max(k, 16)]
for sc, rscore in zip(pool, reranker.scores(query, [s.chunk.text for s in pool])):
sc.score = float(rscore)
pool.sort(key=lambda s: s.score, reverse=True)
return pool[:k]
return scored[:k]
|