Spaces:
Sleeping
Sleeping
File size: 1,215 Bytes
f051f2e |
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 |
\
from typing import List, Tuple
from sentence_transformers import SentenceTransformer
import numpy as np
import faiss
class SessionRAG:
def __init__(self, model_name="sentence-transformers/all-MiniLM-L6-v2"):
self.model = SentenceTransformer(model_name)
self.docs: List[Tuple[str, str]] = []
self.index = None
self.vecs = None
def add_docs(self, items: List[Tuple[str, str]]):
self.docs.extend(items)
texts = [t for _, t in self.docs]
if not texts:
self.index = None; self.vecs=None; return
embs = self.model.encode(texts, convert_to_numpy=True, normalize_embeddings=True).astype(np.float32)
self.vecs = embs
self.index = faiss.IndexFlatIP(embs.shape[1])
self.index.add(embs)
def retrieve(self, query: str, k: int = 6) -> List[str]:
if not self.index or self.vecs is None:
return []
q = self.model.encode([query], convert_to_numpy=True, normalize_embeddings=True).astype(np.float32)
D, I = self.index.search(q, k)
out = []
for idx in I[0]:
if 0 <= idx < len(self.docs):
out.append(self.docs[idx][1])
return out
|