|
|
|
|
|
from __future__ import annotations |
|
|
import json, logging, os |
|
|
from pathlib import Path |
|
|
from typing import List, Dict, Optional |
|
|
import numpy as np |
|
|
import faiss |
|
|
from sentence_transformers import SentenceTransformer |
|
|
|
|
|
log = logging.getLogger(__name__) |
|
|
|
|
|
class Retriever: |
|
|
def __init__(self, kb_path: str = "data/kb.jsonl", |
|
|
model_name: str = "sentence-transformers/all-MiniLM-L6-v2", |
|
|
top_k: int = 4): |
|
|
self.kb_path = Path(kb_path) |
|
|
self.top_k = top_k |
|
|
if not self.kb_path.exists(): |
|
|
raise FileNotFoundError(f"KB file not found: {self.kb_path} (jsonl with {{text,source}})") |
|
|
|
|
|
|
|
|
cache_dir = Path(os.getenv("HF_HOME", "./.cache")) |
|
|
cache_dir.mkdir(parents=True, exist_ok=True) |
|
|
|
|
|
self.model = SentenceTransformer(model_name, cache_folder=str(cache_dir)) |
|
|
|
|
|
self.docs: List[Dict[str, str]] = [] |
|
|
with self.kb_path.open("r", encoding="utf-8") as f: |
|
|
for line in f: |
|
|
line = line.strip() |
|
|
if not line: |
|
|
continue |
|
|
self.docs.append(json.loads(line)) |
|
|
texts = [d["text"] for d in self.docs] |
|
|
emb = self.model.encode(texts, convert_to_numpy=True, normalize_embeddings=True, show_progress_bar=False) |
|
|
self.dim = int(emb.shape[1]) |
|
|
self.index = faiss.IndexFlatIP(self.dim) |
|
|
self.index.add(emb.astype("float32")) |
|
|
|
|
|
def retrieve(self, query: str, k: Optional[int] = None) -> List[Dict]: |
|
|
k = k or self.top_k |
|
|
vec = self.model.encode([query], convert_to_numpy=True, normalize_embeddings=True) |
|
|
D, I = self.index.search(vec.astype("float32"), k) |
|
|
out: List[Dict] = [] |
|
|
for idx, score in zip(I[0], D[0]): |
|
|
if int(idx) < 0: |
|
|
continue |
|
|
d = self.docs[int(idx)] |
|
|
out.append({"text": d["text"], "source": d.get("source", f"kb:{idx}"), "score": float(score)}) |
|
|
return out |
|
|
|
