LawBot / src /apps /utils /retriever.py
Vishwanath77's picture
Upload 2 files
535ee95 verified
import os
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" # Temporary fix
os.environ["FAISS_NO_OPENMP"] = "1" # Prevent FAISS from using OpenMP
import faiss
from numpy.linalg import norm
BASE_DIR = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
def get_path(folder, filename):
# Try the standard 'volumes/' path
path1 = os.path.join(BASE_DIR, "volumes", folder, filename)
if os.path.exists(path1):
return path1
# Try the root-level path
path2 = os.path.join(BASE_DIR, folder, filename)
if os.path.exists(path2):
return path2
return path1
index_path = get_path("indexes", "law_corpus_index2.bin")
if not os.path.exists(index_path):
print(f"⚠️ ERROR: Index file not found at {index_path}")
# Fail gracefully or provide a placeholder for the index if needed
# (Leaving it to crash on faiss.read_index if it truly doesn't exist)
pass
index = faiss.read_index(index_path)
print("Index loaded successfully!")
print("Number of vectors in the index:", index.ntotal)
def vector_db_retriever(query_embeddings, top_k=10):
query_embeddings = query_embeddings / norm(query_embeddings[0])
distances, indices = index.search(query_embeddings, top_k)
return indices, distances