File size: 2,762 Bytes
f35583f | 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 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 | import numpy as np
try:
import faiss
from sentence_transformers import SentenceTransformer
except Exception as e:
print(f"VectorService: Failed to import dependencies: {e}")
faiss = None
SentenceTransformer = None
class VectorService:
def __init__(self):
self.model = None
self.index = None
self.chunks = []
if SentenceTransformer:
# Load model once. This might be slow on startup.
print("Loading generic embedding model (all-MiniLM-L6-v2)...")
try:
self.model = SentenceTransformer("all-MiniLM-L6-v2")
print("Embedding model loaded successfully.")
except Exception as e:
print(f"Failed to load embedding model: {e}")
def create_index_from_results(self, results: list):
"""
Takes a list of search result dicts, creates embeddings, and builds a FAISS index.
"""
if not self.model or not faiss:
print("VectorService: Dependencies missing or model not loaded.")
return
self.chunks = []
texts_to_embed = []
for res in results:
# Combine Title and Content for a rich embedding context
text = f"Title: {res.get('title', '')}\nContent: {res.get('content', '')}"
self.chunks.append(res) # Keep reference to original object
texts_to_embed.append(text)
if not texts_to_embed:
return
try:
embeddings = self.model.encode(texts_to_embed)
dimension = embeddings.shape[1]
self.index = faiss.IndexFlatL2(dimension)
self.index.add(np.array(embeddings))
print(f"VectorService: Created FAISS index with {self.index.ntotal} vectors")
except Exception as e:
print(f"VectorService Error during indexing: {e}")
def search_similar(self, query: str, k: int = 3):
"""
Searches the FAISS index for the most relevant chunks to the query.
"""
if not self.index or not self.model:
return []
try:
query_emb = self.model.encode([query])
distances, indices = self.index.search(query_emb, k)
top_results = []
for idx in indices[0]:
if idx < len(self.chunks) and idx >= 0:
top_results.append(self.chunks[idx])
return top_results
except Exception as e:
print(f"VectorService Error during search: {e}")
return []
vector_service = VectorService()
|