Spaces:
Sleeping
Sleeping
File size: 1,214 Bytes
14fa872 c044be1 14fa872 |
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 |
import logging
try:
import faiss
_HAS_FAISS = True
except ImportError:
logging.warning("FAISS not installed — retrieval will be disabled. Install faiss-cpu or faiss-gpu for full functionality.")
_HAS_FAISS = False
from sentence_transformers import SentenceTransformer
# load embedding model (still works even if FAISS missing)
_model = SentenceTransformer("all-MiniLM-L6-v2")
_index = None
_docs = []
def init_retriever(docs=None):
"""
Initialize FAISS index if FAISS is available.
docs: list[str] to index
"""
global _index, _docs
if not _HAS_FAISS:
_docs = docs or []
return
if docs:
_docs = docs
embeddings = _model.encode(docs, convert_to_numpy=True)
d = embeddings.shape[1]
_index = faiss.IndexFlatL2(d)
_index.add(embeddings)
def retrieve_context(query: str, k: int = 5):
"""
Retrieve top-k docs matching query.
Falls back to empty list if FAISS unavailable.
"""
if not _HAS_FAISS or _index is None or not _docs:
return []
q_emb = _model.encode([query], convert_to_numpy=True)
D, I = _index.search(q_emb, k)
return [_docs[i] for i in I[0] if i < len(_docs)]
|